Merge pull request #726 from narumiruna/rebalance/valuemap

rebalance: replace float slice by string-value map

config/rebalance.yaml

@@ -18,5 +18,4 @@ exchangeStrategies:
       threshold: 2%
       # max amount to buy or sell per order
       maxAmount: 10_000
-      verbose: true
       dryRun: false

pkg/fixedpoint/value_map.go

@@ -0,0 +1,135 @@
+package fixedpoint
+
+type ValueMap map[string]Value
+
+func (m ValueMap) Eq(n ValueMap) bool {
+	if len(m) != len(n) {
+		return false
+	}
+
+	for m_k, m_v := range m {
+		n_v, ok := n[m_k]
+		if !ok {
+			return false
+		}
+
+		if !m_v.Eq(n_v) {
+			return false
+		}
+	}
+
+	return true
+}
+
+func (m ValueMap) Add(n ValueMap) ValueMap {
+	if len(m) != len(n) {
+		panic("unequal length")
+	}
+
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Add(n[k])
+	}
+
+	return o
+}
+
+func (m ValueMap) Sub(n ValueMap) ValueMap {
+	if len(m) != len(n) {
+		panic("unequal length")
+	}
+
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Sub(n[k])
+	}
+
+	return o
+}
+
+func (m ValueMap) Mul(n ValueMap) ValueMap {
+	if len(m) != len(n) {
+		panic("unequal length")
+	}
+
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Mul(n[k])
+	}
+
+	return o
+}
+
+func (m ValueMap) Div(n ValueMap) ValueMap {
+	if len(m) != len(n) {
+		panic("unequal length")
+	}
+
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Div(n[k])
+	}
+
+	return o
+}
+
+func (m ValueMap) AddScalar(x Value) ValueMap {
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Add(x)
+	}
+
+	return o
+}
+
+func (m ValueMap) SubScalar(x Value) ValueMap {
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Sub(x)
+	}
+
+	return o
+}
+
+func (m ValueMap) MulScalar(x Value) ValueMap {
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Mul(x)
+	}
+
+	return o
+}
+
+func (m ValueMap) DivScalar(x Value) ValueMap {
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Div(x)
+	}
+
+	return o
+}
+
+func (m ValueMap) Sum() Value {
+	var sum Value
+	for _, v := range m {
+		sum = sum.Add(v)
+	}
+	return sum
+}
+
+func (m ValueMap) Normalize() ValueMap {
+	sum := m.Sum()
+	if sum.Eq(Zero) {
+		panic("zero sum")
+	}
+
+	return m.DivScalar(sum)
+}

pkg/fixedpoint/value_map_test.go

@@ -0,0 +1,123 @@
+package fixedpoint
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_ValueMap_Eq(t *testing.T) {
+	m1 := ValueMap{
+		"A": NewFromFloat(3.0),
+		"B": NewFromFloat(4.0),
+	}
+
+	m2 := ValueMap{}
+
+	m3 := ValueMap{"A": NewFromFloat(5.0)}
+
+	m4 := ValueMap{
+		"A": NewFromFloat(6.0),
+		"B": NewFromFloat(7.0),
+	}
+
+	m5 := ValueMap{
+		"A": NewFromFloat(3.0),
+		"B": NewFromFloat(4.0),
+	}
+
+	assert.True(t, m1.Eq(m1))
+	assert.False(t, m1.Eq(m2))
+	assert.False(t, m1.Eq(m3))
+	assert.False(t, m1.Eq(m4))
+	assert.True(t, m1.Eq(m5))
+}
+
+func Test_ValueMap_Add(t *testing.T) {
+	m1 := ValueMap{
+		"A": NewFromFloat(3.0),
+		"B": NewFromFloat(4.0),
+	}
+
+	m2 := ValueMap{
+		"A": NewFromFloat(5.0),
+		"B": NewFromFloat(6.0),
+	}
+
+	m3 := ValueMap{
+		"A": NewFromFloat(8.0),
+		"B": NewFromFloat(10.0),
+	}
+
+	m4 := ValueMap{"A": NewFromFloat(8.0)}
+
+	assert.Equal(t, m3, m1.Add(m2))
+	assert.Panics(t, func() { m1.Add(m4) })
+}
+
+func Test_ValueMap_AddScalar(t *testing.T) {
+	x := NewFromFloat(5.0)
+
+	m1 := ValueMap{
+		"A": NewFromFloat(3.0),
+		"B": NewFromFloat(4.0),
+	}
+
+	m2 := ValueMap{
+		"A": NewFromFloat(3.0).Add(x),
+		"B": NewFromFloat(4.0).Add(x),
+	}
+
+	assert.Equal(t, m2, m1.AddScalar(x))
+}
+
+func Test_ValueMap_DivScalar(t *testing.T) {
+	x := NewFromFloat(5.0)
+
+	m1 := ValueMap{
+		"A": NewFromFloat(3.0),
+		"B": NewFromFloat(4.0),
+	}
+
+	m2 := ValueMap{
+		"A": NewFromFloat(3.0).Div(x),
+		"B": NewFromFloat(4.0).Div(x),
+	}
+
+	assert.Equal(t, m2, m1.DivScalar(x))
+}
+
+func Test_ValueMap_Sum(t *testing.T) {
+	m := ValueMap{
+		"A": NewFromFloat(3.0),
+		"B": NewFromFloat(4.0),
+	}
+
+	assert.Equal(t, NewFromFloat(7.0), m.Sum())
+}
+
+func Test_ValueMap_Normalize(t *testing.T) {
+	a := NewFromFloat(3.0)
+	b := NewFromFloat(4.0)
+
+	m := ValueMap{
+		"A": a,
+		"B": b,
+	}
+
+	n := ValueMap{
+		"A": a.Div(a.Add(b)),
+		"B": b.Div(a.Add(b)),
+	}
+
+	assert.True(t, m.Normalize().Eq(n))
+}
+
+func Test_ValueMap_Normalize_zero_sum(t *testing.T) {
+	m := ValueMap{
+		"A": Zero,
+		"B": Zero,
+	}
+
+	assert.Panics(t, func() { m.Normalize() })
+}

pkg/strategy/rebalance/strategy.go

@@ -3,8 +3,6 @@ package rebalance
 import (
 	"context"
 	"fmt"
-	"math"
-	"sort"
 
 	"github.com/sirupsen/logrus"
 
@@ -24,33 +22,18 @@ func init() {
 type Strategy struct {
 	Notifiability *bbgo.Notifiability
 
-	Interval      types.Interval              `json:"interval"`
-	BaseCurrency  string                      `json:"baseCurrency"`
-	TargetWeights map[string]fixedpoint.Value `json:"targetWeights"`
-	Threshold     fixedpoint.Value            `json:"threshold"`
-	Verbose       bool                        `json:"verbose"`
-	DryRun        bool                        `json:"dryRun"`
+	Interval      types.Interval      `json:"interval"`
+	BaseCurrency  string              `json:"baseCurrency"`
+	TargetWeights fixedpoint.ValueMap `json:"targetWeights"`
+	Threshold     fixedpoint.Value    `json:"threshold"`
+	DryRun        bool                `json:"dryRun"`
 	// max amount to buy or sell per order
 	MaxAmount fixedpoint.Value `json:"maxAmount"`
 
-	// sorted currencies
-	currencies []string
-
-	// symbol for subscribing kline
-	symbol string
-
 	activeOrderBook *bbgo.ActiveOrderBook
 }
 
 func (s *Strategy) Initialize() error {
-	for currency := range s.TargetWeights {
-		s.currencies = append(s.currencies, currency)
-	}
-
-	sort.Strings(s.currencies)
-
-	s.symbol = s.currencies[0] + s.BaseCurrency
-
 	return nil
 }
 
@@ -63,6 +46,10 @@ func (s *Strategy) Validate() error {
 		return fmt.Errorf("targetWeights should not be empty")
 	}
 
+	if !s.TargetWeights.Sum().Eq(fixedpoint.One) {
+		return fmt.Errorf("the sum of targetWeights should be 1")
+	}
+
 	for currency, weight := range s.TargetWeights {
 		if weight.Float64() < 0 {
 			return fmt.Errorf("%s weight: %f should not less than 0", currency, weight.Float64())
@@ -81,33 +68,33 @@ func (s *Strategy) Validate() error {
 }
 
 func (s *Strategy) Subscribe(session *bbgo.ExchangeSession) {
-	session.Subscribe(types.KLineChannel, s.symbol, types.SubscribeOptions{Interval: s.Interval})
+	session.Subscribe(types.KLineChannel, s.symbols()[0], types.SubscribeOptions{Interval: s.Interval})
 }
 
 func (s *Strategy) Run(ctx context.Context, orderExecutor bbgo.OrderExecutor, session *bbgo.ExchangeSession) error {
 	s.activeOrderBook = bbgo.NewActiveOrderBook("")
 	s.activeOrderBook.BindStream(session.UserDataStream)
 
-	session.MarketDataStream.OnKLineClosed(func(kline types.KLine) {
-		// cancel active orders before rebalance
-		if err := session.Exchange.CancelOrders(ctx, s.activeOrderBook.Orders()...); err != nil {
-			log.WithError(err).Errorf("failed to cancel orders")
+	markets := session.Markets()
+	for _, symbol := range s.symbols() {
+		if _, ok := markets[symbol]; !ok {
+			return fmt.Errorf("exchange: %s does not supoort matket: %s", session.Exchange.Name(), symbol)
 		}
-
-		s.rebalance(ctx, orderExecutor, session)
-	})
-	return nil
-}
-
-func (s *Strategy) rebalance(ctx context.Context, orderExecutor bbgo.OrderExecutor, session *bbgo.ExchangeSession) {
-	prices, err := s.prices(ctx, session)
-	if err != nil {
-		return
 	}
 
-	marketValues := prices.Mul(s.quantities(session))
+	session.MarketDataStream.OnKLineClosed(func(kline types.KLine) {
+		s.rebalance(ctx, orderExecutor, session)
+	})
 
-	submitOrders := s.generateSubmitOrders(prices, marketValues)
+	return nil
+}
+func (s *Strategy) rebalance(ctx context.Context, orderExecutor bbgo.OrderExecutor, session *bbgo.ExchangeSession) {
+	// cancel active orders before rebalance
+	if err := session.Exchange.CancelOrders(ctx, s.activeOrderBook.Orders()...); err != nil {
+		log.WithError(err).Errorf("failed to cancel orders")
+	}
+
+	submitOrders := s.generateSubmitOrders(ctx, session)
 	for _, order := range submitOrders {
 		log.Infof("generated submit order: %s", order.String())
 	}
@@ -125,44 +112,50 @@ func (s *Strategy) rebalance(ctx context.Context, orderExecutor bbgo.OrderExecut
 	s.activeOrderBook.Add(createdOrders...)
 }
 
-func (s *Strategy) prices(ctx context.Context, session *bbgo.ExchangeSession) (prices types.Float64Slice, err error) {
+func (s *Strategy) prices(ctx context.Context, session *bbgo.ExchangeSession) fixedpoint.ValueMap {
+	m := make(fixedpoint.ValueMap)
+
 	tickers, err := session.Exchange.QueryTickers(ctx, s.symbols()...)
 	if err != nil {
-		return nil, err
+		log.WithError(err).Error("failed to query tickers")
+		return nil
 	}
 
-	for _, currency := range s.currencies {
+	for currency := range s.TargetWeights {
 		if currency == s.BaseCurrency {
-			prices = append(prices, 1.0)
+			m[s.BaseCurrency] = fixedpoint.One
 			continue
 		}
-
-		symbol := currency + s.BaseCurrency
-		prices = append(prices, tickers[symbol].Last.Float64())
+		m[currency] = tickers[currency+s.BaseCurrency].Last
 	}
-	return prices, nil
+
+	return m
 }
 
-func (s *Strategy) quantities(session *bbgo.ExchangeSession) (quantities types.Float64Slice) {
+func (s *Strategy) quantities(session *bbgo.ExchangeSession) fixedpoint.ValueMap {
+	m := make(fixedpoint.ValueMap)
+
 	balances := session.GetAccount().Balances()
-	for _, currency := range s.currencies {
-		quantities = append(quantities, balances[currency].Total().Float64())
+	for currency := range s.TargetWeights {
+		m[currency] = balances[currency].Total()
 	}
-	return quantities
+
+	return m
 }
 
-func (s *Strategy) generateSubmitOrders(prices, marketValues types.Float64Slice) (submitOrders []types.SubmitOrder) {
+func (s *Strategy) generateSubmitOrders(ctx context.Context, session *bbgo.ExchangeSession) (submitOrders []types.SubmitOrder) {
+	prices := s.prices(ctx, session)
+	marketValues := prices.Mul(s.quantities(session))
 	currentWeights := marketValues.Normalize()
 
-	for i, currency := range s.currencies {
+	for currency, targetWeight := range s.TargetWeights {
 		if currency == s.BaseCurrency {
 			continue
 		}
 
 		symbol := currency + s.BaseCurrency
-		currentWeight := currentWeights[i]
-		currentPrice := prices[i]
-		targetWeight := s.TargetWeights[currency].Float64()
+		currentWeight := currentWeights[currency]
+		currentPrice := prices[currency]
 
 		log.Infof("%s price: %v, current weight: %v, target weight: %v",
 			symbol,
@@ -172,8 +165,8 @@ func (s *Strategy) generateSubmitOrders(prices, marketValues types.Float64Slice)
 
 		// calculate the difference between current weight and target weight
 		// if the difference is less than threshold, then we will not create the order
-		weightDifference := targetWeight - currentWeight
-		if math.Abs(weightDifference) < s.Threshold.Float64() {
+		weightDifference := targetWeight.Sub(currentWeight)
+		if weightDifference.Abs().Compare(s.Threshold) < 0 {
 			log.Infof("%s weight distance |%v - %v| = |%v| less than the threshold: %v",
 				symbol,
 				currentWeight,
@@ -183,7 +176,7 @@ func (s *Strategy) generateSubmitOrders(prices, marketValues types.Float64Slice)
 			continue
 		}
 
-		quantity := fixedpoint.NewFromFloat((weightDifference * marketValues.Sum()) / currentPrice)
+		quantity := weightDifference.Mul(marketValues.Sum()).Div(currentPrice)
 
 		side := types.SideTypeBuy
 		if quantity.Sign() < 0 {
@@ -192,7 +185,7 @@ func (s *Strategy) generateSubmitOrders(prices, marketValues types.Float64Slice)
 		}
 
 		if s.MaxAmount.Sign() > 0 {
-			quantity = bbgo.AdjustQuantityByMaxAmount(quantity, fixedpoint.NewFromFloat(currentPrice), s.MaxAmount)
+			quantity = bbgo.AdjustQuantityByMaxAmount(quantity, currentPrice, s.MaxAmount)
 			log.Infof("adjust the quantity %v (%s %s @ %v) by max amount %v",
 				quantity,
 				symbol,
@@ -200,22 +193,25 @@ func (s *Strategy) generateSubmitOrders(prices, marketValues types.Float64Slice)
 				currentPrice,
 				s.MaxAmount)
 		}
+
 		log.Debugf("symbol: %v, quantity: %v", symbol, quantity)
+
 		order := types.SubmitOrder{
 			Symbol:   symbol,
 			Side:     side,
 			Type:     types.OrderTypeLimit,
 			Quantity: quantity,
-			Price:    fixedpoint.NewFromFloat(currentPrice),
+			Price:    currentPrice,
 		}
 
 		submitOrders = append(submitOrders, order)
 	}
+
 	return submitOrders
 }
 
 func (s *Strategy) symbols() (symbols []string) {
-	for _, currency := range s.currencies {
+	for currency := range s.TargetWeights {
 		if currency == s.BaseCurrency {
 			continue
 		}