energy stats algorithm and test

Author: zhouselena

internal/cmd/perfcomp/energystatistics.go

@@ -5,3 +5,42 @@
 // a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
 
 package main
+
+import (
+	"math"
+
+	"gonum.org/v1/gonum/mat"
+)
+
+// Given two matrices, this function returns
+// (e, t, h) = (E-statistic, test statistic, e-coefficient of inhomogeneity)
+func getEnergyStatistics(x, y *mat.Dense) (float64, float64, float64) {
+	n, _ := x.Dims()
+	m, _ := y.Dims()
+	nf := float64(n)
+	mf := float64(m)
+	A := getDistance(x, y) / (nf * mf) // E|X-Y|
+	B := getDistance(x, x) / (nf * nf) // E|X-X'|
+	C := getDistance(y, y) / (mf * mf) // E|Y-Y'|
+
+	E := 2*A - B - C // D^2(F_x, F_y)
+	T := ((nf * mf) / (nf + mf)) * E
+	H := E / (2 * A)
+	return E, T, H
+}
+
+// Given two vectors (expected 1 col),
+// this function returns the sum of distances between each pair.
+func getDistance(x, y *mat.Dense) float64 {
+	xrows, _ := x.Dims()
+	yrows, _ := y.Dims()
+
+	var sum float64
+
+	for i := 0; i < xrows; i++ {
+		for j := 0; j < yrows; j++ {
+			sum += math.Sqrt(math.Pow((x.At(i, 0) - y.At(j, 0)), 2))
+		}
+	}
+	return sum
+}

internal/cmd/perfcomp/energystatistics_test.go

@@ -5,3 +5,65 @@
 // a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
 
 package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"gonum.org/v1/gonum/mat"
+)
+
+func createTestVectors(start1 int, stop1 int, step1 int, start2 int, stop2 int, step2 int) (*mat.Dense, *mat.Dense) {
+	xData := []float64{}
+	i := start1
+	for i < stop1 {
+		xData = append(xData, float64(i))
+		i += step1
+	}
+
+	yData := []float64{}
+	j := start2
+	for j < stop2 {
+		yData = append(yData, float64(j))
+		j += step2
+	}
+
+	x := mat.NewDense(len(xData), 1, xData)
+	y := mat.NewDense(len(yData), 1, yData)
+
+	return x, y
+}
+
+func TestEnergyStatistics(t *testing.T) {
+
+	t.Run("similar distributions should have small e,t,h values ", func(t *testing.T) {
+		x, y := createTestVectors(1, 100, 1, 1, 105, 1)
+		e, tstat, h := getEnergyStatistics(x, y)
+
+		del := 1e-3
+
+		assert.InDelta(t, 0.160, e, del) // |0.160 - e| < 1/100
+		assert.InDelta(t, 8.136, tstat, del)
+		assert.InDelta(t, 0.002, h, del)
+	})
+
+	t.Run("different distributions should have large e,t,h values", func(t *testing.T) {
+		x, y := createTestVectors(1, 100, 1, 10000, 13000, 14)
+		e, tstat, h := getEnergyStatistics(x, y)
+		del := 1e-3
+
+		assert.InDelta(t, 21859.691, e, del)
+		assert.InDelta(t, 1481794.709, tstat, del)
+		assert.InDelta(t, 0.954, h, del)
+	})
+
+	t.Run("uni-variate distributions", func(t *testing.T) {
+		x, y := createTestVectors(1, 300, 1, 1000, 5000, 10)
+		e, tstat, h := getEnergyStatistics(x, y)
+		del := 1e-3
+
+		assert.InDelta(t, 4257.009, e, del)
+		assert.InDelta(t, 1481794.709, tstat, del)
+		assert.InDelta(t, 0.954, h, del)
+	})
+}