diff --git a/formula/wma.go b/formula/wma.go new file mode 100644 index 0000000000000000000000000000000000000000..4e0aae67cf29307ff0af4fa5cefcc2fa9aac71c7 --- /dev/null +++ b/formula/wma.go @@ -0,0 +1,36 @@ +package formula + +import ( + "gitee.com/quant1x/pandas" + "gitee.com/quant1x/pandas/exception" + "gitee.com/quant1x/pandas/stat" +) + +// WMA 通达信S序列的N日加权移动平均 Yn = (1*X1+2*X2+3*X3+...+n*Xn)/(1+2+3+...+Xn) +func WMA(S pandas.Series, N any) any { + var X []float32 + switch v := N.(type) { + case int: + X = stat.Repeat[float32](float32(v), S.Len()) + case pandas.Series: + vs := v.Values() + X = pandas.SliceToFloat32(vs) + X = stat.Align(X, pandas.Nil2Float32, S.Len()) + default: + panic(exception.New(1, "error window")) + } + return S.Rolling(X).Apply(func(S pandas.Series, N float32) float32 { + if S.Len() == 0 { + return stat.Nil2Float32 + } + x := pandas.ToFloat32(S) + //fmt.Println(x) + x = stat.Reverse(x) + //fmt.Println(x) + v := stat.CumSum(x) + //fmt.Println(v) + v1 := stat.Sum(v) + v2 := v1 * 2 / N / (N + 1) + return v2 + }).Values() +} diff --git a/formula/wma_test.go b/formula/wma_test.go new file mode 100644 index 0000000000000000000000000000000000000000..5b8e36dbb3127e644b7ce07d46d7eb727eee4222 --- /dev/null +++ b/formula/wma_test.go @@ -0,0 +1,20 @@ +package formula + +import ( + "fmt" + "gitee.com/quant1x/pandas" + "testing" +) + +func TestWMA(t *testing.T) { + f0 := []float64{1, 2, 3, 4} + //f1 := []float64{70, 80, 75, 83, 86} + //f2 := []float64{90, 69, 60, 88, 87} + + s0 := pandas.NewSeriesWithoutType("f0", f0) + //s1 := pandas.NewSeriesWithoutType("f1", f1) + //s2 := pandas.NewSeriesWithoutType("f2", f2) + fmt.Println(WMA(s0, 4)) + //fmt.Println(WMA(s1, 5)) + //fmt.Println(WMA(s2, 5)) +} diff --git a/generic_rolling.go b/generic_rolling.go index 251c81f1309ec39af7b69b87e120ef152c31cb16..25ebb7b657ca7f128eaf64cb49377a3e472580e8 100644 --- a/generic_rolling.go +++ b/generic_rolling.go @@ -5,14 +5,6 @@ import ( "gitee.com/quant1x/pandas/stat" ) -// RollingV1 滑动窗口 -func (self *NDFrame) RollingV1(window int) RollingWindowV1 { - return RollingWindowV1{ - window: window, - series: self, - } -} - // RollingAndExpandingMixin 滚动和扩展静态横切 type RollingAndExpandingMixin struct { window []float32 @@ -56,3 +48,13 @@ func (r RollingAndExpandingMixin) getBlocks() (blocks []Series) { return } + +// Apply 接受一个回调 +func (r RollingAndExpandingMixin) Apply(f func(S Series, N float32) float32) (s Series) { + s = r.series.Empty() + for i, block := range r.getBlocks() { + v := f(block, r.window[i]) + s.Append(v) + } + return +} diff --git a/rolling_v1.go b/rolling_v1.go index cb10d224f7c479c82047e9c3cccb403799bde4c1..ad88d374d6277797912e9606bda9d31b9834b94b 100644 --- a/rolling_v1.go +++ b/rolling_v1.go @@ -6,6 +6,14 @@ type RollingWindowV1 struct { series Series } +// RollingV1 滑动窗口 +func (self *NDFrame) RollingV1(window int) RollingWindowV1 { + return RollingWindowV1{ + window: window, + series: self, + } +} + func (r RollingWindowV1) getBlocks() (blocks []Series) { for i := 1; i <= r.series.Len(); i++ { if i < r.window { diff --git a/stat/cumsum.go b/stat/cumsum.go new file mode 100644 index 0000000000000000000000000000000000000000..2879a2dbd95f3e6f8cc2cb789cc2ab9cae2fa374 --- /dev/null +++ b/stat/cumsum.go @@ -0,0 +1,35 @@ +package stat + +import ( + "github.com/viterin/vek" + "github.com/viterin/vek/vek32" + "golang.org/x/exp/slices" + "unsafe" +) + +// CumSum 计算累和 +func CumSum[T StatType](f []T) []T { + if len(f) == 0 { + return []T{} + } + var d any + var s any + s = f + bitSize := unsafe.Sizeof(f[0]) + if bitSize == 4 { + d = vek32.CumSum(s.([]float32)) + } else if bitSize == 8 { + d = vek.CumSum(s.([]float64)) + } else { + // 剩下的就是int32和int64, 循环吧 + sum := T(0) + x := slices.Clone(f) + for i := 0; i < len(x); i++ { + sum += x[i] + x[i] = sum + } + d = x + } + + return d.([]T) +} diff --git a/stat/reverse.go b/stat/reverse.go new file mode 100644 index 0000000000000000000000000000000000000000..d73d39a10977ed2f8fb717dbd7772c8ac9f5d8b6 --- /dev/null +++ b/stat/reverse.go @@ -0,0 +1,9 @@ +package stat + +// Reverse 反转切片 +func Reverse[S ~[]E, E any](s S) S { + for i, j := 0, len(s)-1; i < j; i, j = i+1, j-1 { + s[i], s[j] = s[j], s[i] + } + return s +} diff --git a/stat/reverse_test.go b/stat/reverse_test.go new file mode 100644 index 0000000000000000000000000000000000000000..7439f171da12b58ce20dce4a6316e3975b58c03b --- /dev/null +++ b/stat/reverse_test.go @@ -0,0 +1,18 @@ +package stat + +import ( + "fmt" + "testing" +) + +func TestReverse(t *testing.T) { + f0 := []float64{1.1, 2.2, 1.3, 1.4} + f1 := []float64{70, 80, 75, 83, 86} + f2 := []float64{90, 69, 60, 88, 87} + r0 := f0[0:1:1] + fmt.Println(Reverse(f0)) + fmt.Println(Reverse(f1)) + fmt.Println(Reverse(f2)) + + _ = r0 +}