ENH: add 'fweights' and 'aweights' arguments to covariance calculations.

'fweights' allows integer frequencies to be specified for observation vectors,
and 'aweights' provides a more general importance or probabalistic weighting.

1 files changed, 72 insertions, 5 deletions

diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index 12f9d414b..ad71fd3fa 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -1391,9 +1391,20 @@ class TestCorrCoef(TestCase):
 
 
 class TestCov(TestCase):
+    x1 = np.array([[0, 2], [1, 1], [2, 0]]).T
+    res1 = np.array([[1., -1.], [-1., 1.]])
+    x2 = np.array([0.0, 1.0, 2.0], ndmin=2)
+    frequencies = np.array([1, 4, 1])
+    x2_repeats = np.array([[0.0], [1.0], [1.0], [1.0], [1.0], [2.0]]).T
+    res2 = np.array([[0.4, -0.4], [-0.4, 0.4]])
+    unit_frequencies = np.ones(3, dtype=np.integer)
+    weights = np.array([1.0, 4.0, 1.0])
+    res3 = np.array([[2./3., -2./3.], [-2./3., 2./3.]])
+    unit_weights = np.ones(3)
+    x3 = np.array([0.3942, 0.5969, 0.7730, 0.9918, 0.7964])
+
     def test_basic(self):
-        x = np.array([[0, 2], [1, 1], [2, 0]]).T
-        assert_allclose(cov(x), np.array([[1., -1.], [-1., 1.]]))
+        assert_allclose(cov(self.x1), self.res1)
 
     def test_complex(self):
         x = np.array([[1, 2, 3], [1j, 2j, 3j]])
@@ -1414,11 +1425,67 @@ class TestCov(TestCase):
                                np.array([[np.nan, np.nan], [np.nan, np.nan]]))
 
     def test_wrong_ddof(self):
-        x = np.array([[0, 2], [1, 1], [2, 0]]).T
         with warnings.catch_warnings(record=True):
             warnings.simplefilter('always', RuntimeWarning)
-            assert_array_equal(cov(x, ddof=5),
-                               np.array([[np.inf, -np.inf], [-np.inf, np.inf]]))
+            assert_array_equal(cov(self.x1, ddof=5),
+                               np.array([[np.inf, -np.inf],
+                                         [-np.inf, np.inf]]))
+
+    def test_1D_rowvar(self):
+        assert_allclose(cov(self.x3), cov(self.x3, rowvar=0))
+        y = np.array([0.0780, 0.3107, 0.2111, 0.0334, 0.8501])
+        assert_allclose(cov(self.x3, y), cov(self.x3, y, rowvar=0))
+
+    def test_1D_variance(self):
+        assert_allclose(cov(self.x3, ddof=1), np.var(self.x3, ddof=1))
+
+    def test_fweights(self):
+        assert_allclose(cov(self.x2, fweights=self.frequencies),
+                        cov(self.x2_repeats))
+        assert_allclose(cov(self.x1, fweights=self.frequencies),
+                        self.res2)
+        assert_allclose(cov(self.x1, fweights=self.unit_frequencies),
+                        self.res1)
+        nonint = self.frequencies + 0.5
+        assert_raises(TypeError, cov, self.x1, fweights=nonint)
+        f = np.ones((2, 3), dtype=np.integer)
+        assert_raises(RuntimeError, cov, self.x1, fweights=f)
+        f = np.ones(2, dtype=np.integer)
+        assert_raises(RuntimeError, cov, self.x1, fweights=f)
+        f = -1*np.ones(3, dtype=np.integer)
+        assert_raises(ValueError, cov, self.x1, fweights=f)
+
+    def test_aweights(self):
+        assert_allclose(cov(self.x1, aweights=self.weights), self.res3)
+        assert_allclose(cov(self.x1, aweights=3.0*self.weights),
+                        cov(self.x1, aweights=self.weights))
+        assert_allclose(cov(self.x1, aweights=self.unit_weights), self.res1)
+        w = np.ones((2, 3))
+        assert_raises(RuntimeError, cov, self.x1, aweights=w)
+        w = np.ones(2)
+        assert_raises(RuntimeError, cov, self.x1, aweights=w)
+        w = -1.0*np.ones(3)
+        assert_raises(ValueError, cov, self.x1, aweights=w)
+
+    def test_unit_fweights_and_aweights(self):
+        assert_allclose(cov(self.x2, fweights=self.frequencies,
+                            aweights=self.unit_weights),
+                        cov(self.x2_repeats))
+        assert_allclose(cov(self.x1, fweights=self.frequencies,
+                            aweights=self.unit_weights),
+                        self.res2)
+        assert_allclose(cov(self.x1, fweights=self.unit_frequencies,
+                            aweights=self.unit_weights),
+                        self.res1)
+        assert_allclose(cov(self.x1, fweights=self.unit_frequencies,
+                            aweights=self.weights),
+                        self.res3)
+        assert_allclose(cov(self.x1, fweights=self.unit_frequencies,
+                            aweights=3.0*self.weights),
+                        cov(self.x1, aweights=self.weights))
+        assert_allclose(cov(self.x1, fweights=self.unit_frequencies,
+                            aweights=self.unit_weights),
+                        self.res1)
 
 
 class Test_I0(TestCase):