From 1f58387cb6188b11d1719fa43d55f22a038f08fe Mon Sep 17 00:00:00 2001
From: Brendan Tracey <tracey.brendan@gmail.com>
Date: Fri, 29 Sep 2017 07:03:31 -0600
Subject: [PATCH] stat: add BivariateMoment (#244)

* stat: add MixedMoment and fix bug in Moment code

* Revert unrelated changes, rename to BivariateMoment

* last fix
---
 stat/stat.go      | 35 +++++++++++++++++++++++++++++
 stat/stat_test.go | 57 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+)

diff --git a/stat/stat.go b/stat/stat.go
index 11b90e91..8696e91c 100644
--- a/stat/stat.go
+++ b/stat/stat.go
@@ -862,6 +862,41 @@ func Mode(x, weights []float64) (val float64, count float64) {
 	return max, maxCount
 }
 
+// BivariateMoment computes the weighted mixed moment between the samples x and y.
+//  E[(x - μ_x)^r*(y - μ_y)^s]
+// No degrees of freedom correction is done.
+// The lengths of x and y must be equal. If weights is nil then all of the
+// weights are 1. If weights is not nil, then len(x) must equal len(weights).
+func BivariateMoment(r, s float64, x, y, weights []float64) float64 {
+	meanX := Mean(x, weights)
+	meanY := Mean(y, weights)
+	if len(x) != len(y) {
+		panic("stat: slice length mismatch")
+	}
+	if weights == nil {
+		var m float64
+		for i, vx := range x {
+			vy := y[i]
+			m += math.Pow(vx-meanX, r) * math.Pow(vy-meanY, s)
+		}
+		return m / float64(len(x))
+	}
+	if len(weights) != len(x) {
+		panic("stat: slice length mismatch")
+	}
+	var (
+		m          float64
+		sumWeights float64
+	)
+	for i, vx := range x {
+		vy := y[i]
+		w := weights[i]
+		m += w * math.Pow(vx-meanX, r) * math.Pow(vy-meanY, s)
+		sumWeights += w
+	}
+	return m / sumWeights
+}
+
 // Moment computes the weighted n^th moment of the samples,
 //  E[(x - μ)^N]
 // No degrees of freedom correction is done.
diff --git a/stat/stat_test.go b/stat/stat_test.go
index bdb24dd3..b8fd9adc 100644
--- a/stat/stat_test.go
+++ b/stat/stat_test.go
@@ -1094,6 +1094,63 @@ func TestMode(t *testing.T) {
 	}
 }
 
+func TestMixedMoment(t *testing.T) {
+	for i, test := range []struct {
+		x, y, weights []float64
+		r, s          float64
+		ans           float64
+	}{
+		{
+			x:   []float64{10, 2, 1, 8, 5},
+			y:   []float64{8, 15, 1, 6, 3},
+			r:   1,
+			s:   1,
+			ans: 0.48,
+		},
+		{
+			x:       []float64{10, 2, 1, 8, 5},
+			y:       []float64{8, 15, 1, 6, 3},
+			weights: []float64{1, 1, 1, 1, 1},
+			r:       1,
+			s:       1,
+			ans:     0.48,
+		},
+		{
+			x:       []float64{10, 2, 1, 8, 5},
+			y:       []float64{8, 15, 1, 6, 3},
+			weights: []float64{2, 3, 0.2, 8, 4},
+			r:       1,
+			s:       1,
+			ans:     -4.786371011357490,
+		},
+		{
+			x:       []float64{10, 2, 1, 8, 5},
+			y:       []float64{8, 15, 1, 6, 3},
+			weights: []float64{2, 3, 0.2, 8, 4},
+			r:       2,
+			s:       3,
+			ans:     1.598600579313326e+03,
+		},
+	} {
+		m := BivariateMoment(test.r, test.s, test.x, test.y, test.weights)
+		if math.Abs(test.ans-m) > 1e-14 {
+			t.Errorf("Moment mismatch case %d. Expected %v, found %v", i, test.ans, m)
+		}
+	}
+	if !Panics(func() { BivariateMoment(1, 1, make([]float64, 3), make([]float64, 2), nil) }) {
+		t.Errorf("Moment did not panic with x, y length mismatch")
+	}
+	if !Panics(func() { BivariateMoment(1, 1, make([]float64, 2), make([]float64, 3), nil) }) {
+		t.Errorf("Moment did not panic with x, y length mismatch")
+	}
+	if !Panics(func() { BivariateMoment(1, 1, make([]float64, 2), make([]float64, 2), make([]float64, 3)) }) {
+		t.Errorf("Moment did not panic with x, weights length mismatch")
+	}
+	if !Panics(func() { BivariateMoment(1, 1, make([]float64, 2), make([]float64, 2), make([]float64, 1)) }) {
+		t.Errorf("Moment did not panic with x, weights length mismatch")
+	}
+}
+
 func TestMoment(t *testing.T) {
 	for i, test := range []struct {
 		x       []float64