grid2: add grid struct

pkg/strategy/grid2/grid.go

@@ -0,0 +1,102 @@
+package grid2
+
+import (
+	"github.com/c9s/bbgo/pkg/fixedpoint"
+)
+
+type Grid struct {
+	UpperPrice fixedpoint.Value `json:"upperPrice"`
+	LowerPrice fixedpoint.Value `json:"lowerPrice"`
+
+	// Spread is the spread of each grid
+	Spread fixedpoint.Value `json:"spread"`
+
+	// Size is the number of total grids
+	Size fixedpoint.Value `json:"size"`
+
+	// Pins are the pinned grid prices, from low to high
+	Pins []fixedpoint.Value `json:"pins"`
+
+	pinsCache map[fixedpoint.Value]struct{} `json:"-"`
+}
+
+func NewGrid(lower, upper, density fixedpoint.Value) *Grid {
+	var height = upper - lower
+	var size = height.Div(density)
+	var pins []fixedpoint.Value
+
+	for p := lower; p <= upper; p += size {
+		pins = append(pins, p)
+	}
+
+	grid := &Grid{
+		UpperPrice: upper,
+		LowerPrice: lower,
+		Size:       density,
+		Spread:     size,
+		Pins:       pins,
+		pinsCache:  make(map[fixedpoint.Value]struct{}, len(pins)),
+	}
+	grid.updatePinsCache()
+	return grid
+}
+
+func (g *Grid) Above(price fixedpoint.Value) bool {
+	return price > g.UpperPrice
+}
+
+func (g *Grid) Below(price fixedpoint.Value) bool {
+	return price < g.LowerPrice
+}
+
+func (g *Grid) OutOfRange(price fixedpoint.Value) bool {
+	return price < g.LowerPrice || price > g.UpperPrice
+}
+
+func (g *Grid) updatePinsCache() {
+	for _, pin := range g.Pins {
+		g.pinsCache[pin] = struct{}{}
+	}
+}
+
+func (g *Grid) HasPin(pin fixedpoint.Value) (ok bool) {
+	_, ok = g.pinsCache[pin]
+	return ok
+}
+
+func (g *Grid) ExtendUpperPrice(upper fixedpoint.Value) (newPins []fixedpoint.Value) {
+	g.UpperPrice = upper
+
+	// since the grid is extended, the size should be updated as well
+	g.Size = (g.UpperPrice - g.LowerPrice).Div(g.Spread).Floor()
+
+	lastPin := g.Pins[len(g.Pins)-1]
+	for p := lastPin + g.Spread; p <= g.UpperPrice; p += g.Spread {
+		newPins = append(newPins, p)
+	}
+
+	g.Pins = append(g.Pins, newPins...)
+	g.updatePinsCache()
+	return newPins
+}
+
+func (g *Grid) ExtendLowerPrice(lower fixedpoint.Value) (newPins []fixedpoint.Value) {
+	g.LowerPrice = lower
+
+	// since the grid is extended, the size should be updated as well
+	g.Size = (g.UpperPrice - g.LowerPrice).Div(g.Spread).Floor()
+
+	firstPin := g.Pins[0]
+	numToAdd := (firstPin - g.LowerPrice).Div(g.Spread).Floor()
+	if numToAdd == 0 {
+		return newPins
+	}
+
+	for p := firstPin - g.Spread.Mul(numToAdd); p < firstPin; p += g.Spread {
+		newPins = append(newPins, p)
+	}
+
+	g.Pins = append(newPins, g.Pins...)
+	g.updatePinsCache()
+	return newPins
+}

pkg/strategy/grid2/grid_test.go

@@ -0,0 +1,81 @@
+package grid2
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/c9s/bbgo/pkg/fixedpoint"
+)
+
+func TestNewGrid(t *testing.T) {
+	upper := fixedpoint.NewFromFloat(500.0)
+	lower := fixedpoint.NewFromFloat(100.0)
+	size := fixedpoint.NewFromFloat(100.0)
+	grid := NewGrid(lower, upper, size)
+	assert.Equal(t, upper, grid.UpperPrice)
+	assert.Equal(t, lower, grid.LowerPrice)
+	assert.Equal(t, fixedpoint.NewFromFloat(4), grid.Spread)
+	if assert.Len(t, grid.Pins, 101) {
+		assert.Equal(t, fixedpoint.NewFromFloat(100.0), grid.Pins[0])
+		assert.Equal(t, fixedpoint.NewFromFloat(500.0), grid.Pins[100])
+	}
+}
+
+func TestGrid_HasPin(t *testing.T) {
+	upper := fixedpoint.NewFromFloat(500.0)
+	lower := fixedpoint.NewFromFloat(100.0)
+	size := fixedpoint.NewFromFloat(100.0)
+	grid := NewGrid(lower, upper, size)
+
+	assert.True(t, grid.HasPin(fixedpoint.NewFromFloat(100.0)))
+	assert.True(t, grid.HasPin(fixedpoint.NewFromFloat(500.0)))
+	assert.False(t, grid.HasPin(fixedpoint.NewFromFloat(101.0)))
+}
+
+func TestGrid_ExtendUpperPrice(t *testing.T) {
+	upper := fixedpoint.NewFromFloat(500.0)
+	lower := fixedpoint.NewFromFloat(100.0)
+	size := fixedpoint.NewFromFloat(100.0)
+	grid := NewGrid(lower, upper, size)
+
+	originalSpread := grid.Spread
+	newPins := grid.ExtendUpperPrice(fixedpoint.NewFromFloat(1000.0))
+	assert.Equal(t, originalSpread, grid.Spread)
+	assert.Len(t, newPins, 125) // (1000-500) / 4
+	assert.Equal(t, fixedpoint.NewFromFloat(4), grid.Spread)
+	if assert.Len(t, grid.Pins, 226) {
+		assert.Equal(t, fixedpoint.NewFromFloat(100.0), grid.Pins[0])
+		assert.Equal(t, fixedpoint.NewFromFloat(1000.0), grid.Pins[225])
+	}
+}
+
+func TestGrid_ExtendLowerPrice(t *testing.T) {
+	upper := fixedpoint.NewFromFloat(3000.0)
+	lower := fixedpoint.NewFromFloat(2000.0)
+	size := fixedpoint.NewFromFloat(100.0)
+	grid := NewGrid(lower, upper, size)
+
+	// spread = (3000 - 2000) / 100.0
+	expectedSpread := fixedpoint.NewFromFloat(10.0)
+	assert.Equal(t, expectedSpread, grid.Spread)
+
+	originalSpread := grid.Spread
+	newPins := grid.ExtendLowerPrice(fixedpoint.NewFromFloat(1000.0))
+	assert.Equal(t, originalSpread, grid.Spread)
+
+	// 100 = (2000-1000) / 10
+	if assert.Len(t, newPins, 100) {
+		assert.Equal(t, fixedpoint.NewFromFloat(2000.0)-expectedSpread, newPins[99])
+	}
+
+	assert.Equal(t, expectedSpread, grid.Spread)
+	if assert.Len(t, grid.Pins, 201) {
+		assert.Equal(t, fixedpoint.NewFromFloat(1000.0), grid.Pins[0])
+		assert.Equal(t, fixedpoint.NewFromFloat(3000.0), grid.Pins[200])
+	}
+
+	newPins2 := grid.ExtendLowerPrice(
+		fixedpoint.NewFromFloat(1000.0 - 1.0))
+	assert.Len(t, newPins2, 0) // should have no new pin generated
+}