feature: add correlation for series

pkg/types/indicator.go

@@ -643,4 +643,94 @@ func Stdev(a Series, length int) float64 {
 	return math.Sqrt(s / float64(length))
 }
 
+type CorrFunc func(Series, Series, int) float64
+
+func Kendall(a, b Series, length int) float64 {
+	if a.Length() < length {
+		length = a.Length()
+	}
+	if b.Length() < length {
+		length = b.Length()
+	}
+	aRanks := Rank(a, length)
+	bRanks := Rank(b, length)
+	concordant, discordant := 0, 0
+	for i := 0; i < length; i++ {
+		for j := i + 1; j < length; j++ {
+			value := (aRanks.Index(i) - aRanks.Index(j)) * (bRanks.Index(i) - bRanks.Index(j))
+			if value > 0 {
+				concordant++
+			} else {
+				discordant++
+			}
+		}
+	}
+	return float64(concordant-discordant) * 2.0 / float64(length*(length-1))
+}
+
+func Rank(a Series, length int) Series {
+	if length > a.Length() {
+		length = a.Length()
+	}
+	rank := make([]float64, length)
+	mapper := make([]float64, length+1)
+	for i := length - 1; i >= 0; i-- {
+		ii := a.Index(i)
+		counter := 0.
+		for j := 0; j < length; j++ {
+			if a.Index(j) <= ii {
+				counter += 1.
+			}
+		}
+		rank[i] = counter
+		mapper[int(counter)] += 1.
+	}
+	output := NewQueue(length)
+	for i := length - 1; i >= 0; i-- {
+		output.Update(rank[i] - (mapper[int(rank[i])]-1.)/2)
+	}
+	return output
+}
+
+func Pearson(a, b Series, length int) float64 {
+	if a.Length() < length {
+		length = a.Length()
+	}
+	if b.Length() < length {
+		length = b.Length()
+	}
+	x := make([]float64, length)
+	y := make([]float64, length)
+	for i := 0; i < length; i++ {
+		x[i] = a.Index(i)
+		y[i] = b.Index(i)
+	}
+	return stat.Correlation(x, y, nil)
+}
+
+func Spearman(a, b Series, length int) float64 {
+	if a.Length() < length {
+		length = a.Length()
+	}
+	if b.Length() < length {
+		length = b.Length()
+	}
+	aRank := Rank(a, length)
+	bRank := Rank(b, length)
+	return Pearson(aRank, bRank, length)
+}
+
+// similar to pandas.Series.corr() function.
+//
+// method could either be `types.Pearson`, `types.Spearman` or `types.Kendall`
+func Correlation(a Series, b Series, length int, method ...CorrFunc) float64 {
+	var runner CorrFunc
+	if len(method) == 0 {
+		runner = Pearson
+	} else {
+		runner = method[0]
+	}
+	return runner(a, b, length)
+}
+
 // TODO: ta.linreg

pkg/types/indicator_test.go

@@ -33,3 +33,26 @@ func TestFloat64Slice(t *testing.T) {
 	b = append(b, 3.0)
 	assert.Equal(t, c.Last(), 1.)
 }
+
+/*
+python
+
+import pandas as pd
+s1 = pd.Series([.2, 0., .6, .2, .2])
+s2 = pd.Series([.3, .6, .0, .1])
+print(s1.corr(s2, method='pearson'))
+print(s1.corr(s2, method='spearman')
+print(s1.corr(s2, method='kendall'))
+print(s1.rank())
+*/
+func TestCorr(t *testing.T) {
+	var a = Float64Slice{.2, .0, .6, .2}
+	var b = Float64Slice{.3, .6, .0, .1}
+	corr := Correlation(&a, &b, 4, Pearson)
+	assert.InDelta(t, corr, -0.8510644, 0.001)
+	out := Rank(&a, 4)
+	assert.Equal(t, out.Index(0), 2.5)
+	assert.Equal(t, out.Index(1), 4.0)
+	corr = Correlation(&a, &b, 4, Spearman)
+	assert.InDelta(t, corr, -0.94868, 0.001)
+}