indicator: refactor ATRStream

This commit is contained in:
c9s 2023-05-30 12:29:50 +08:00
parent da15f47f17
commit f65d6267fc
No known key found for this signature in database
GPG Key ID: 7385E7E464CB0A54
2 changed files with 38 additions and 22 deletions

View File

@ -1,36 +1,55 @@
package indicator package indicator
import ( import (
"math"
"github.com/c9s/bbgo/pkg/types" "github.com/c9s/bbgo/pkg/types"
) )
type ATRStream struct { type ATRStream struct {
// embedded struct
Float64Series Float64Series
types.SeriesBase // parameters
types.IntervalWindow
window int // private states
multiplier float64 rma *RMAStream
window int
previousClose float64
} }
func ATR2(source KLineSubscription, window int) *ATRStream { func ATR2(source KLineSubscription, window int) *ATRStream {
s := &ATRStream{ s := &ATRStream{
Float64Series: NewFloat64Series(), Float64Series: NewFloat64Series(),
window: window, window: window,
multiplier: 2.0 / float64(1+window),
} }
s.rma = RMA2(s, window, true)
source.AddSubscriber(func(k types.KLine) { source.AddSubscriber(func(k types.KLine) {
// v := s.mapper(k) s.calculateAndPush(k.High.Float64(), k.Low.Float64(), k.Close.Float64())
// s.slice.Push(v)
// s.EmitUpdate(v)
}) })
return s return s
} }
func (s *ATRStream) calculateAndPush(k types.KLine) { func (s *ATRStream) calculateAndPush(high, low, cls float64) {
// v2 := s.calculate(v) if s.previousClose == .0 {
// s.slice.Push(v2) s.previousClose = cls
// s.EmitUpdate(v2) return
}
trueRange := high - low
hc := math.Abs(high - s.previousClose)
lc := math.Abs(low - s.previousClose)
if trueRange < hc {
trueRange = hc
}
if trueRange < lc {
trueRange = lc
}
s.previousClose = cls
s.slice.Push(trueRange)
s.rma.EmitUpdate(trueRange)
} }

View File

@ -1,25 +1,22 @@
package indicator package indicator
import (
"github.com/c9s/bbgo/pkg/types"
)
type RMAStream struct { type RMAStream struct {
// embedded structs // embedded structs
Float64Series Float64Series
// config fields // config fields
types.IntervalWindow
Adjust bool Adjust bool
window int
counter int counter int
sum, previous float64 sum, previous float64
} }
func RMA2(source Float64Source, iw types.IntervalWindow) *RMAStream { func RMA2(source Float64Source, window int, adjust bool) *RMAStream {
s := &RMAStream{ s := &RMAStream{
Float64Series: NewFloat64Series(), Float64Series: NewFloat64Series(),
IntervalWindow: iw, window: window,
Adjust: adjust,
} }
if sub, ok := source.(Float64Subscription); ok { if sub, ok := source.(Float64Subscription); ok {
@ -39,7 +36,7 @@ func (s *RMAStream) calculateAndPush(v float64) {
} }
func (s *RMAStream) calculate(x float64) float64 { func (s *RMAStream) calculate(x float64) float64 {
lambda := 1 / float64(s.Window) lambda := 1 / float64(s.window)
tmp := 0.0 tmp := 0.0
if s.counter == 0 { if s.counter == 0 {
s.sum = 1 s.sum = 1
@ -54,7 +51,7 @@ func (s *RMAStream) calculate(x float64) float64 {
} }
s.counter++ s.counter++
if s.counter < s.Window { if s.counter < s.window {
// we can use x, but we need to use 0. to make the same behavior as the result from python pandas_ta // we can use x, but we need to use 0. to make the same behavior as the result from python pandas_ta
s.slice.Push(0) s.slice.Push(0)
} }