indicator: rewrite VWMA calculator

c9s · Jul 14, 2022 · dd3bd6a · dd3bd6a
1 parent 2ef8ecf
commit dd3bd6a
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 34 deletions.
diff --git a/pkg/indicator/inf.go b/pkg/indicator/inf.go
@@ -6,6 +6,12 @@ type KLineWindowUpdater interface {
 	OnKLineWindowUpdate(func(interval types.Interval, window types.KLineWindow))
 }
 
+type KLineClosedBinder interface {
+	BindK(target KLineClosedEmitter, symbol string, interval types.Interval)
+}
+
+// KLineClosedEmitter is currently applied to the market data stream
+// the market data stream emits the KLine closed event to the listeners.
 type KLineClosedEmitter interface {
 	OnKLineClosed(func(k types.KLine))
 }

diff --git a/pkg/indicator/sma.go b/pkg/indicator/sma.go
@@ -86,6 +86,10 @@ func (inc *SMA) handleKLineWindowUpdate(interval types.Interval, window types.KL
 	inc.CalculateAndUpdate(window)
 }
 
+func (inc *SMA) BindK(target KLineClosedEmitter, symbol string, interval types.Interval) {
+	target.OnKLineClosed(types.KLineWith(symbol, interval, inc.PushK))
+}
+
 func (inc *SMA) Bind(updater KLineWindowUpdater) {
 	updater.OnKLineWindowUpdate(inc.handleKLineWindowUpdate)
 }

diff --git a/pkg/indicator/vwma.go b/pkg/indicator/vwma.go
@@ -3,8 +3,6 @@ package indicator
 import (
 	"time"
 
-	log "github.com/sirupsen/logrus"
-
 	"github.com/c9s/bbgo/pkg/types"
 )
 
@@ -22,10 +20,14 @@ Volume Weighted Moving Average
 type VWMA struct {
 	types.SeriesBase
 	types.IntervalWindow
-	Values  types.Float64Slice
+
+	Values         types.Float64Slice
+	PriceVolumeSMA *SMA
+	VolumeSMA      *SMA
+
 	EndTime time.Time
 
-	UpdateCallbacks []func(value float64)
+	updateCallbacks []func(value float64)
 }
 
 func (inc *VWMA) Last() float64 {
@@ -49,46 +51,46 @@ func (inc *VWMA) Length() int {
 
 var _ types.SeriesExtend = &VWMA{}
 
-
-func (inc *VWMA) CalculateAndUpdate(allKLines []types.KLine) {
-	if len(allKLines) < inc.Window {
-		return
+func (inc *VWMA) Update(price, volume float64) {
+	if inc.PriceVolumeSMA == nil {
+		inc.PriceVolumeSMA = &SMA{IntervalWindow: inc.IntervalWindow}
+		inc.SeriesBase.Series = inc
 	}
 
-	var index = len(allKLines) - 1
-	var kline = allKLines[index]
-
-	if inc.EndTime != zeroTime && kline.EndTime.Before(inc.EndTime) {
-		return
+	if inc.VolumeSMA == nil {
+		inc.VolumeSMA = &SMA{IntervalWindow: inc.IntervalWindow}
 	}
 
-	var recentK = allKLines[index-(inc.Window-1) : index+1]
+	inc.PriceVolumeSMA.Update(price * volume)
+	inc.VolumeSMA.Update(volume)
 
-	pv, err := calculateSMA(recentK, inc.Window, KLinePriceVolumeMapper)
-	if err != nil {
-		log.WithError(err).Error("price x volume SMA error")
-		return
-	}
-	v, err := calculateSMA(recentK, inc.Window, KLineVolumeMapper)
-	if err != nil {
-		log.WithError(err).Error("volume SMA error")
+	pv := inc.PriceVolumeSMA.Last()
+	v := inc.VolumeSMA.Last()
+	vwma := pv / v
+	inc.Values.Push(vwma)
+}
+
+func (inc *VWMA) CalculateAndUpdate(allKLines []types.KLine) {
+	if len(allKLines) < inc.Window {
 		return
 	}
 
-	if len(inc.Values) == 0 {
-		inc.SeriesBase.Series = inc
-	}
+	var last = allKLines[len(allKLines)-1]
 
-	vwma := pv / v
-	inc.Values.Push(vwma)
+	if inc.VolumeSMA == nil {
+		for _, k := range allKLines {
+			if inc.EndTime != zeroTime && k.EndTime.Before(inc.EndTime) {
+				return
+			}
 
-	if len(inc.Values) > MaxNumOfSMA {
-		inc.Values = inc.Values[MaxNumOfSMATruncateSize-1:]
+			inc.Update(k.Close.Float64(), k.Volume.Float64())
+		}
+	} else {
+		inc.Update(last.Close.Float64(), last.Volume.Float64())
 	}
 
-	inc.EndTime = allKLines[index].EndTime.Time()
-
-	inc.EmitUpdate(vwma)
+	inc.EndTime = last.EndTime.Time()
+	inc.EmitUpdate(inc.Values.Last())
 }
 
 func (inc *VWMA) handleKLineWindowUpdate(interval types.Interval, window types.KLineWindow) {

diff --git a/pkg/indicator/vwma_callbacks.go b/pkg/indicator/vwma_callbacks.go