TST: Add tests for NA support in the polynomial fitting functions.

6 files changed, 175 insertions, 0 deletions

diff --git a/numpy/polynomial/tests/test_chebyshev.py b/numpy/polynomial/tests/test_chebyshev.py
index 64b52fa5d..5e6d5632b 100644
--- a/numpy/polynomial/tests/test_chebyshev.py
+++ b/numpy/polynomial/tests/test_chebyshev.py
@@ -436,6 +436,35 @@ class TestFitting(TestCase):
         wcoef2d = cheb.chebfit(x, np.array([yw,yw]).T, 3, w=w)
         assert_almost_equal(wcoef2d, np.array([coef3,coef3]).T)
 
+        #test NA
+        y = f(x)
+        y[10] = 100
+
+        xm = x.view(maskna=1)
+        xm[10] = np.NA
+        res = cheb.chebfit(xm, y, 3)
+        assert_almost_equal(res, coef3)
+
+        ym = y.view(maskna=1)
+        ym[10] = np.NA
+        res = cheb.chebfit(x, ym, 3)
+        assert_almost_equal(res, coef3)
+
+        y2 = np.vstack((y,y)).T
+        y2[10,0] = 100
+        y2[15,1] = 100
+        y2m = y2.view(maskna=1)
+        y2m[10,0] = np.NA
+        y2m[15,1] = np.NA
+        res = cheb.chebfit(x, y2m, 3).T
+        assert_almost_equal(res[0], coef3)
+        assert_almost_equal(res[1], coef3)
+
+        wm = np.ones_like(x, maskna=1)
+        wm[10] = np.NA
+        res = cheb.chebfit(x, y, 3, w=wm)
+        assert_almost_equal(res, coef3)
+
 
 class TestGauss(TestCase):
 
diff --git a/numpy/polynomial/tests/test_hermite.py b/numpy/polynomial/tests/test_hermite.py
index 576f5dde7..425ebe5e0 100644
--- a/numpy/polynomial/tests/test_hermite.py
+++ b/numpy/polynomial/tests/test_hermite.py
@@ -426,6 +426,35 @@ class TestFitting(TestCase):
         wcoef2d = herm.hermfit(x, np.array([yw,yw]).T, 3, w=w)
         assert_almost_equal(wcoef2d, np.array([coef3,coef3]).T)
 
+        #test NA
+        y = f(x)
+        y[10] = 100
+
+        xm = x.view(maskna=1)
+        xm[10] = np.NA
+        res = herm.hermfit(xm, y, 3)
+        assert_almost_equal(res, coef3)
+
+        ym = y.view(maskna=1)
+        ym[10] = np.NA
+        res = herm.hermfit(x, ym, 3)
+        assert_almost_equal(res, coef3)
+
+        y2 = np.vstack((y,y)).T
+        y2[10,0] = 100
+        y2[15,1] = 100
+        y2m = y2.view(maskna=1)
+        y2m[10,0] = np.NA
+        y2m[15,1] = np.NA
+        res = herm.hermfit(x, y2m, 3).T
+        assert_almost_equal(res[0], coef3)
+        assert_almost_equal(res[1], coef3)
+
+        wm = np.ones_like(x, maskna=1)
+        wm[10] = np.NA
+        res = herm.hermfit(x, y, 3, w=wm)
+        assert_almost_equal(res, coef3)
+
 
 class TestGauss(TestCase):
 
diff --git a/numpy/polynomial/tests/test_hermite_e.py b/numpy/polynomial/tests/test_hermite_e.py
index c7b5a7b17..f0ecb66b8 100644
--- a/numpy/polynomial/tests/test_hermite_e.py
+++ b/numpy/polynomial/tests/test_hermite_e.py
@@ -423,6 +423,35 @@ class TestFitting(TestCase):
         wcoef2d = herme.hermefit(x, np.array([yw,yw]).T, 3, w=w)
         assert_almost_equal(wcoef2d, np.array([coef3,coef3]).T)
 
+        #test NA
+        y = f(x)
+        y[10] = 100
+
+        xm = x.view(maskna=1)
+        xm[10] = np.NA
+        res = herme.hermefit(xm, y, 3)
+        assert_almost_equal(res, coef3)
+
+        ym = y.view(maskna=1)
+        ym[10] = np.NA
+        res = herme.hermefit(x, ym, 3)
+        assert_almost_equal(res, coef3)
+
+        y2 = np.vstack((y,y)).T
+        y2[10,0] = 100
+        y2[15,1] = 100
+        y2m = y2.view(maskna=1)
+        y2m[10,0] = np.NA
+        y2m[15,1] = np.NA
+        res = herme.hermefit(x, y2m, 3).T
+        assert_almost_equal(res[0], coef3)
+        assert_almost_equal(res[1], coef3)
+
+        wm = np.ones_like(x, maskna=1)
+        wm[10] = np.NA
+        res = herme.hermefit(x, y, 3, w=wm)
+        assert_almost_equal(res, coef3)
+
 
 class TestGauss(TestCase):
 
diff --git a/numpy/polynomial/tests/test_laguerre.py b/numpy/polynomial/tests/test_laguerre.py
index 1d3e0442f..8b3fd5f8a 100644
--- a/numpy/polynomial/tests/test_laguerre.py
+++ b/numpy/polynomial/tests/test_laguerre.py
@@ -421,6 +421,35 @@ class TestFitting(TestCase):
         wcoef2d = lag.lagfit(x, np.array([yw,yw]).T, 3, w=w)
         assert_almost_equal(wcoef2d, np.array([coef3,coef3]).T)
 
+        #test NA
+        y = f(x)
+        y[10] = 100
+
+        xm = x.view(maskna=1)
+        xm[10] = np.NA
+        res = lag.lagfit(xm, y, 3)
+        assert_almost_equal(res, coef3)
+
+        ym = y.view(maskna=1)
+        ym[10] = np.NA
+        res = lag.lagfit(x, ym, 3)
+        assert_almost_equal(res, coef3)
+
+        y2 = np.vstack((y,y)).T
+        y2[10,0] = 100
+        y2[15,1] = 100
+        y2m = y2.view(maskna=1)
+        y2m[10,0] = np.NA
+        y2m[15,1] = np.NA
+        res = lag.lagfit(x, y2m, 3).T
+        assert_almost_equal(res[0], coef3)
+        assert_almost_equal(res[1], coef3)
+
+        wm = np.ones_like(x, maskna=1)
+        wm[10] = np.NA
+        res = lag.lagfit(x, y, 3, w=wm)
+        assert_almost_equal(res, coef3)
+
 
 class TestGauss(TestCase):
 
diff --git a/numpy/polynomial/tests/test_legendre.py b/numpy/polynomial/tests/test_legendre.py
index e00d2777c..c7b71fc05 100644
--- a/numpy/polynomial/tests/test_legendre.py
+++ b/numpy/polynomial/tests/test_legendre.py
@@ -424,6 +424,35 @@ class TestFitting(TestCase):
         wcoef2d = leg.legfit(x, np.array([yw,yw]).T, 3, w=w)
         assert_almost_equal(wcoef2d, np.array([coef3,coef3]).T)
 
+        #test NA
+        y = f(x)
+        y[10] = 100
+
+        xm = x.view(maskna=1)
+        xm[10] = np.NA
+        res = leg.legfit(xm, y, 3)
+        assert_almost_equal(res, coef3)
+
+        ym = y.view(maskna=1)
+        ym[10] = np.NA
+        res = leg.legfit(x, ym, 3)
+        assert_almost_equal(res, coef3)
+
+        y2 = np.vstack((y,y)).T
+        y2[10,0] = 100
+        y2[15,1] = 100
+        y2m = y2.view(maskna=1)
+        y2m[10,0] = np.NA
+        y2m[15,1] = np.NA
+        res = leg.legfit(x, y2m, 3).T
+        assert_almost_equal(res[0], coef3)
+        assert_almost_equal(res[1], coef3)
+
+        wm = np.ones_like(x, maskna=1)
+        wm[10] = np.NA
+        res = leg.legfit(x, y, 3, w=wm)
+        assert_almost_equal(res, coef3)
+
 
 class TestGauss(TestCase):
 
diff --git a/numpy/polynomial/tests/test_polynomial.py b/numpy/polynomial/tests/test_polynomial.py
index bae711cbf..4854f95ee 100644
--- a/numpy/polynomial/tests/test_polynomial.py
+++ b/numpy/polynomial/tests/test_polynomial.py
@@ -441,6 +441,36 @@ class TestMisc(TestCase) :
         wcoef2d = poly.polyfit(x, np.array([yw,yw]).T, 3, w=w)
         assert_almost_equal(wcoef2d, np.array([coef3,coef3]).T)
 
+        #test NA
+        y = f(x)
+        y[10] = 100
+
+        xm = x.view(maskna=1)
+        xm[10] = np.NA
+        res = poly.polyfit(xm, y, 3)
+        assert_almost_equal(res, coef3)
+
+        ym = y.view(maskna=1)
+        ym[10] = np.NA
+        res = poly.polyfit(x, ym, 3)
+        assert_almost_equal(res, coef3)
+
+        y2 = np.vstack((y,y)).T
+        y2[10,0] = 100
+        y2[15,1] = 100
+        y2m = y2.view(maskna=1)
+        y2m[10,0] = np.NA
+        y2m[15,1] = np.NA
+        res = poly.polyfit(x, y2m, 3).T
+        assert_almost_equal(res[0], coef3)
+        assert_almost_equal(res[1], coef3)
+
+        wm = np.ones_like(x, maskna=1)
+        wm[10] = np.NA
+        res = poly.polyfit(x, y, 3, w=wm)
+        assert_almost_equal(res, coef3)
+
+
     def test_polytrim(self) :
         coef = [2, -1, 1, 0]