bbgo: refactor standard indicator set

pkg/bbgo/standard_indicator_set.go

@@ -9,27 +9,23 @@ import (
 )
 
 var (
-	debugEWMA = false
-	debugSMA  = false
 	debugBOLL = false
 )
 
 func init() {
 	// when using --dotenv option, the dotenv is loaded from command.PersistentPreRunE, not init.
 	// hence here the env var won't enable the debug flag
-	util.SetEnvVarBool("DEBUG_EWMA", &debugEWMA)
-	util.SetEnvVarBool("DEBUG_SMA", &debugSMA)
 	util.SetEnvVarBool("DEBUG_BOLL", &debugBOLL)
 }
 
 type StandardIndicatorSet struct {
 	Symbol string
+
 	// Standard indicators
 	// interval -> window
-	sma   map[types.IntervalWindow]*indicator.SMA
-	ewma  map[types.IntervalWindow]*indicator.EWMA
 	boll    map[types.IntervalWindowBandWidth]*indicator.BOLL
 	stoch   map[types.IntervalWindow]*indicator.STOCH
+	simples map[types.IntervalWindow]indicator.Simple
 
 	stream types.Stream
 	store  *MarketDataStore
@@ -38,105 +34,80 @@ type StandardIndicatorSet struct {
 func NewStandardIndicatorSet(symbol string, stream types.Stream, store *MarketDataStore) *StandardIndicatorSet {
 	return &StandardIndicatorSet{
 		Symbol:  symbol,
-		sma:    make(map[types.IntervalWindow]*indicator.SMA),
-		ewma:   make(map[types.IntervalWindow]*indicator.EWMA),
-		boll:   make(map[types.IntervalWindowBandWidth]*indicator.BOLL),
-		stoch:  make(map[types.IntervalWindow]*indicator.STOCH),
 		store:   store,
 		stream:  stream,
+		simples: make(map[types.IntervalWindow]indicator.Simple),
+
+		boll:  make(map[types.IntervalWindowBandWidth]*indicator.BOLL),
+		stoch: make(map[types.IntervalWindow]*indicator.STOCH),
 	}
 }
 
+func (set *StandardIndicatorSet) initAndBind(inc indicator.KLinePusher, iw types.IntervalWindow) {
+	if klines, ok := set.store.KLinesOfInterval(iw.Interval); ok {
+		for _, k := range *klines {
+			inc.PushK(k)
+		}
+	}
+
+	set.stream.OnKLineClosed(types.KLineWith(set.Symbol, iw.Interval, inc.PushK))
+}
+
+func (set *StandardIndicatorSet) allocateSimpleIndicator(t indicator.Simple, iw types.IntervalWindow) indicator.Simple {
+	inc, ok := set.simples[iw]
+	if ok {
+		return inc
+	}
+
+	inc = t
+	set.initAndBind(inc, iw)
+	set.simples[iw] = inc
+	return t
+}
+
+// SMA is a helper function that returns the simple moving average indicator of the given interval and the window size.
+func (set *StandardIndicatorSet) SMA(iw types.IntervalWindow) *indicator.SMA {
+	inc := set.allocateSimpleIndicator(&indicator.SMA{IntervalWindow: iw}, iw)
+	return inc.(*indicator.SMA)
+}
+
+// EWMA is a helper function that returns the exponential weighed moving average indicator of the given interval and the window size.
+func (set *StandardIndicatorSet) EWMA(iw types.IntervalWindow) *indicator.EWMA {
+	inc := set.allocateSimpleIndicator(&indicator.EWMA{IntervalWindow: iw}, iw)
+	return inc.(*indicator.EWMA)
+}
+
+func (set *StandardIndicatorSet) PivotLow(iw types.IntervalWindow) *indicator.PivotLow {
+	inc := set.allocateSimpleIndicator(&indicator.PivotLow{IntervalWindow: iw}, iw)
+	return inc.(*indicator.PivotLow)
+}
+
+func (set *StandardIndicatorSet) STOCH(iw types.IntervalWindow) *indicator.STOCH {
+	inc, ok := set.stoch[iw]
+	if !ok {
+		inc = &indicator.STOCH{IntervalWindow: iw}
+		set.initAndBind(inc, iw)
+		set.stoch[iw] = inc
+	}
+
+	return inc
+}
+
 // BOLL returns the bollinger band indicator of the given interval, the window and bandwidth
 func (set *StandardIndicatorSet) BOLL(iw types.IntervalWindow, bandWidth float64) *indicator.BOLL {
 	iwb := types.IntervalWindowBandWidth{IntervalWindow: iw, BandWidth: bandWidth}
 	inc, ok := set.boll[iwb]
 	if !ok {
 		inc = &indicator.BOLL{IntervalWindow: iw, K: bandWidth}
-
+		set.initAndBind(inc, iw)
-		if klines, ok := set.store.KLinesOfInterval(iw.Interval); ok {
-			for _, k := range *klines {
-				inc.PushK(k)
-			}
-		}
 
 		if debugBOLL {
 			inc.OnUpdate(func(sma float64, upBand float64, downBand float64) {
 				logrus.Infof("%s BOLL %s: sma=%f up=%f down=%f", set.Symbol, iw.String(), sma, upBand, downBand)
 			})
 		}
-
-		inc.BindK(set.stream, set.Symbol, iw.Interval)
 		set.boll[iwb] = inc
 	}
 
 	return inc
 }
-
-// SMA returns the simple moving average indicator of the given interval and the window size.
-func (set *StandardIndicatorSet) SMA(iw types.IntervalWindow) *indicator.SMA {
-	inc, ok := set.sma[iw]
-	if !ok {
-		inc = &indicator.SMA{IntervalWindow: iw}
-
-		if klines, ok := set.store.KLinesOfInterval(iw.Interval); ok {
-			for _, k := range *klines {
-				inc.PushK(k)
-			}
-		}
-
-		if debugSMA {
-			inc.OnUpdate(func(value float64) {
-				logrus.Infof("%s SMA %s: %f", set.Symbol, iw.String(), value)
-			})
-		}
-
-		inc.BindK(set.stream, set.Symbol, iw.Interval)
-		set.sma[iw] = inc
-	}
-
-	return inc
-}
-
-// EWMA returns the exponential weighed moving average indicator of the given interval and the window size.
-func (set *StandardIndicatorSet) EWMA(iw types.IntervalWindow) *indicator.EWMA {
-	inc, ok := set.ewma[iw]
-	if !ok {
-		inc = &indicator.EWMA{IntervalWindow: iw}
-
-		if klines, ok := set.store.KLinesOfInterval(iw.Interval); ok {
-			for _, k := range *klines {
-				inc.PushK(k)
-			}
-		}
-
-		if debugEWMA {
-			inc.OnUpdate(func(value float64) {
-				logrus.Infof("%s EWMA %s: value=%f", set.Symbol, iw.String(), value)
-			})
-		}
-
-		inc.BindK(set.stream, set.Symbol, iw.Interval)
-		set.ewma[iw] = inc
-	}
-
-	return inc
-}
-
-func (set *StandardIndicatorSet) STOCH(iw types.IntervalWindow) *indicator.STOCH {
-	inc, ok := set.stoch[iw]
-	if !ok {
-		inc = &indicator.STOCH{IntervalWindow: iw}
-
-		if klines, ok := set.store.KLinesOfInterval(iw.Interval); ok {
-			for _, k := range *klines {
-				inc.PushK(k)
-			}
-		}
-
-		inc.BindK(set.stream, set.Symbol, iw.Interval)
-		set.stoch[iw] = inc
-	}
-
-	return inc
-}

pkg/indicator/interface.go

@@ -22,6 +22,13 @@ type KLinePusher interface {
 	PushK(k types.KLine)
 }
 
+// Simple is the simple indicator that only returns one float64 value
+type Simple interface {
+	KLinePusher
+	Last() float64
+	OnUpdate(f func(value float64))
+}
+
 type KLineCalculateUpdater interface {
 	CalculateAndUpdate(allKLines []types.KLine)
 }

pkg/indicator/pivot_low.go

@@ -53,6 +53,7 @@ func (inc *PivotLow) PushK(k types.KLine) {
 	inc.EmitUpdate(inc.Last())
 }
 
+
 func calculatePivotLow(lows types.Float64Slice, window int) (float64, error) {
 	length := len(lows)
 	if length == 0 || length < window {

pkg/indicator/stoch.go

@@ -69,16 +69,6 @@ func (inc *STOCH) PushK(k types.KLine) {
 	inc.EmitUpdate(inc.LastK(), inc.LastD())
 }
 
-func (inc *STOCH) BindK(target KLineClosedEmitter, symbol string, interval types.Interval) {
-	target.OnKLineClosed(types.KLineWith(symbol, interval, inc.PushK))
-}
-
-func (inc *STOCH) LoadK(allKLines []types.KLine) {
-	for _, k := range allKLines {
-		inc.PushK(k)
-	}
-}
-
 func (inc *STOCH) GetD() types.Series {
 	return &inc.D
 }