qbtrade/pkg/types/pca.go
2024-06-27 22:42:38 +08:00

60 lines
1.4 KiB
Go

package types
import (
"fmt"
"gonum.org/v1/gonum/mat"
)
type PCA struct {
svd *mat.SVD
}
func (pca *PCA) FitTransform(x []SeriesExtend, lookback, feature int) ([]SeriesExtend, error) {
if err := pca.Fit(x, lookback); err != nil {
return nil, err
}
return pca.Transform(x, lookback, feature), nil
}
func (pca *PCA) Fit(x []SeriesExtend, lookback int) error {
vec := make([]float64, lookback*len(x))
for i, xx := range x {
mean := xx.Mean(lookback)
for j := 0; j < lookback; j++ {
vec[i+j*i] = xx.Last(j) - mean
}
}
pca.svd = &mat.SVD{}
diffMatrix := mat.NewDense(lookback, len(x), vec)
if ok := pca.svd.Factorize(diffMatrix, mat.SVDThin); !ok {
return fmt.Errorf("Unable to factorize")
}
return nil
}
func (pca *PCA) Transform(x []SeriesExtend, lookback int, features int) (result []SeriesExtend) {
result = make([]SeriesExtend, features)
vTemp := new(mat.Dense)
pca.svd.VTo(vTemp)
var ret mat.Dense
vec := make([]float64, lookback*len(x))
for i, xx := range x {
for j := 0; j < lookback; j++ {
vec[i+j*i] = xx.Last(j)
}
}
newX := mat.NewDense(lookback, len(x), vec)
ret.Mul(newX, vTemp)
newMatrix := mat.NewDense(lookback, features, nil)
newMatrix.Copy(&ret)
for i := 0; i < features; i++ {
queue := NewQueue(lookback)
for j := 0; j < lookback; j++ {
queue.Update(newMatrix.At(lookback-j-1, i))
}
result[i] = queue
}
return result
}