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
import (
"math"
"github.com/c9s/bbgo/pkg/types"
)
type ATRStream struct {
// embedded struct
Float64Series
types.SeriesBase
// parameters
types.IntervalWindow
window int
multiplier float64
// private states
rma *RMAStream
window int
previousClose float64
}
func ATR2(source KLineSubscription, window int) *ATRStream {
s := &ATRStream{
Float64Series: NewFloat64Series(),
window: window,
multiplier: 2.0 / float64(1+window),
}
s.rma = RMA2(s, window, true)
source.AddSubscriber(func(k types.KLine) {
// v := s.mapper(k)
// s.slice.Push(v)
// s.EmitUpdate(v)
s.calculateAndPush(k.High.Float64(), k.Low.Float64(), k.Close.Float64())
})
return s
}
func (s *ATRStream) calculateAndPush(k types.KLine) {
// v2 := s.calculate(v)
// s.slice.Push(v2)
// s.EmitUpdate(v2)
func (s *ATRStream) calculateAndPush(high, low, cls float64) {
if s.previousClose == .0 {
s.previousClose = cls
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
import (
"github.com/c9s/bbgo/pkg/types"
)
type RMAStream struct {
// embedded structs
Float64Series
// config fields
types.IntervalWindow
Adjust bool
window int
counter int
sum, previous float64
}
func RMA2(source Float64Source, iw types.IntervalWindow) *RMAStream {
func RMA2(source Float64Source, window int, adjust bool) *RMAStream {
s := &RMAStream{
Float64Series: NewFloat64Series(),
IntervalWindow: iw,
Float64Series: NewFloat64Series(),
window: window,
Adjust: adjust,
}
if sub, ok := source.(Float64Subscription); ok {
@ -39,7 +36,7 @@ func (s *RMAStream) calculateAndPush(v float64) {
}
func (s *RMAStream) calculate(x float64) float64 {
lambda := 1 / float64(s.Window)
lambda := 1 / float64(s.window)
tmp := 0.0
if s.counter == 0 {
s.sum = 1
@ -54,7 +51,7 @@ func (s *RMAStream) calculate(x float64) float64 {
}
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
s.slice.Push(0)
}