addition for MeanAndVariance

dlwh · May 16, 2015 · d5ec11e · d5ec11e
1 parent 377ee67
commit d5ec11e
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 4 deletions.
diff --git a/math/src/main/scala/breeze/stats/DescriptiveStats.scala b/math/src/main/scala/breeze/stats/DescriptiveStats.scala
@@ -31,6 +31,21 @@ import spire.implicits.cfor
 
 case class MeanAndVariance(mean: Double, variance: Double, count: Long) {
   def stdDev: Double = math.sqrt(variance)
+
+  def +(other: MeanAndVariance): MeanAndVariance = {
+    val d = other.mean - this.mean
+    val newMean = this.mean + d * other.count / (other.count + this.count)
+
+    val m2a = this.variance * (this.count - 1)
+    val m2b = other.variance * (other.count - 1)
+
+    val m2x = m2a + m2b + d * d * (other.count * this.count)/ (other.count + this.count)
+
+    val newVariance = m2x / (other.count + this.count - 1)
+
+    MeanAndVariance(newMean, newVariance, this.count + other.count)
+
+  }
 }
 
 object accumulateAndCount extends UFunc {

diff --git a/math/src/test/scala/breeze/stats/DescriptiveStatsTest.scala b/math/src/test/scala/breeze/stats/DescriptiveStatsTest.scala
@@ -69,6 +69,7 @@ class DescriptiveStatsTest extends WordSpec with Matchers {
       assert( math.abs( result(0,0) - 1.0) < 1e-7)
       assert( math.abs( result(1,1) - 1.0) < 1e-7)
     }
+
   }
 }
 
@@ -92,10 +93,19 @@ class DescriptiveStatsTest2 extends FunSuite {
     assert(mav == mav2)
   }
 
-  //  test("complex mean") {
-  //    val data =  DenseVector[Complex]( (0.0 + 1.0 * bmath.i), (1.0 + 0.0 * bmath.i), (2.0 + 2.0 * bmath.i) )
-  //    assert( mean(data) == (1.0 + 1.0 * bmath.i), "complex mean incorrect")
-  //  }
+  test("mean and variance addition") {
+    val r = new Random(0)
+    val data =  Array.fill(100000)(r.nextGaussian)
+    val data2 =  Array.fill(100000)(r.nextGaussian * 5 + 3)
+    val mav = meanAndVariance(data)
+    val mav2 = meanAndVariance(data2)
+    val mavTotal = meanAndVariance(data ++ data2)
+    val mavSum = mav + mav2
+    assert(breeze.numerics.closeTo(mavTotal.mean, mavSum.mean, 1E-5))
+    assert(breeze.numerics.closeTo(mavTotal.variance, mavSum.variance, 1E-5))
+    assert(mavSum.count == mavTotal.count)
+  }
+
 
   test("median") {
     val dataOdd =  DenseVector(0,1,2,3,400000)