ENH: missingdata: Change things to help scipy pass its tests

Improving the NumPy tests a bit to catch these errors better...

7 files changed, 134 insertions, 102 deletions

diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py
index f102f4b58..c7e390c73 100644
--- a/numpy/core/fromnumeric.py
+++ b/numpy/core/fromnumeric.py
@@ -1455,7 +1455,14 @@ def sum(a, axis=None, dtype=None, out=None):
             out[...] = res
             return out
         return res
-    return um.add.reduce(a, axis=axis, dtype=dtype, out=out)
+    elif not (type(a) is mu.ndarray):
+        try:
+            sum = a.sum
+        except AttributeError:
+            return um.add.reduce(a, axis=axis, dtype=dtype, out=out)
+        return sum(axis=axis, dtype=dtype, out=out)
+    else:
+        return um.add.reduce(a, axis=axis, dtype=dtype, out=out)
 
 def product (a, axis=None, dtype=None, out=None):
     """
@@ -1998,7 +2005,14 @@ def prod(a, axis=None, dtype=None, out=None):
     True
 
     """
-    return um.multiply.reduce(a, axis=axis, dtype=dtype, out=out)
+    if not (type(a) is mu.ndarray):
+        try:
+            prod = a.prod
+        except AttributeError:
+            return um.multiply.reduce(a, axis=axis, dtype=dtype, out=out)
+        return prod(axis=axis, dtype=dtype, out=out)
+    else:
+        return um.multiply.reduce(a, axis=axis, dtype=dtype, out=out)
 
 def cumprod(a, axis=None, dtype=None, out=None):
     """
diff --git a/numpy/core/src/multiarray/mapping.c b/numpy/core/src/multiarray/mapping.c
index d122eb214..91802ee34 100644
--- a/numpy/core/src/multiarray/mapping.c
+++ b/numpy/core/src/multiarray/mapping.c
@@ -698,8 +698,6 @@ array_boolean_subscript(PyArrayObject *self,
     PyArrayObject *ret;
     int self_has_maskna = PyArray_HASMASKNA(self), needs_api = 0;
     npy_intp bmask_size;
-    int bmask_axes[NPY_MAXDIMS];
-    int *op_axes[2] = {NULL, NULL};
 
     if (PyArray_DESCR(bmask)->type_num != NPY_BOOL) {
         PyErr_SetString(PyExc_TypeError,
@@ -715,30 +713,15 @@ array_boolean_subscript(PyArrayObject *self,
         return NULL;
     }
 
-    /*
-     * If the boolean mask has one dimension, broadcast to
-     * the left instead of to the right. Other broadcasting
-     * is disallowed to minimize inconsistency with NumPy in
-     * general.
-     */
     if (PyArray_NDIM(bmask) != PyArray_NDIM(self)) {
-        int i;
-
-        if (PyArray_NDIM(bmask) != 1) {
-            PyErr_SetString(PyExc_ValueError,
-                    "The boolean mask indexing array "
-                    "is neither one-dimensional nor "
-                    "matches the operand's number of "
-                    "dimensions");
-            return NULL;
-        }
-        op_axes[1] = bmask_axes;
-        bmask_axes[0] = 0;
-        for (i = 1; i < PyArray_NDIM(self); ++i) {
-            bmask_axes[i] = -1;
-        }
+        PyErr_SetString(PyExc_ValueError,
+                "The boolean mask assignment indexing array "
+                "must have the same number of dimensions as "
+                "the array being indexed");
+        return NULL;
     }
 
+
     /*
      * Since we've checked that the mask contains no NAs, we
      * can do a straightforward count of the boolean True values
@@ -803,9 +786,8 @@ array_boolean_subscript(PyArrayObject *self,
          */
         op_flags[1] = NPY_ITER_READONLY | NPY_ITER_IGNORE_MASKNA;
 
-        iter = NpyIter_AdvancedNew(2, op, flags, order, NPY_NO_CASTING,
-                                op_flags, NULL,
-                                PyArray_NDIM(self), op_axes, NULL, 0);
+        iter = NpyIter_MultiNew(2, op, flags, order, NPY_NO_CASTING,
+                                op_flags, NULL);
         if (iter == NULL) {
             Py_DECREF(ret);
             return NULL;
@@ -948,8 +930,6 @@ array_ass_boolean_subscript(PyArrayObject *self,
     int needs_api = 0;
     npy_intp bmask_size;
     char constant_valid_mask = 1;
-    int bmask_axes[NPY_MAXDIMS];
-    int *op_axes[2] = {NULL, NULL};
 
     if (PyArray_DESCR(bmask)->type_num != NPY_BOOL) {
         PyErr_SetString(PyExc_TypeError,
@@ -959,34 +939,19 @@ array_ass_boolean_subscript(PyArrayObject *self,
     }
 
     if (PyArray_NDIM(v) > 1) {
-        PyErr_SetString(PyExc_TypeError,
+        PyErr_Format(PyExc_TypeError,
                 "NumPy boolean array indexing assignment "
-                "requires a 0 or 1-dimensional input");
+                "requires a 0 or 1-dimensional input, input "
+                "has %d dimensions", PyArray_NDIM(v));
         return -1;
     }
 
-    /*
-     * If the boolean mask has one dimension, broadcast to
-     * the left instead of to the right. Other broadcasting
-     * is disallowed to minimize inconsistency with NumPy in
-     * general.
-     */
     if (PyArray_NDIM(bmask) != PyArray_NDIM(self)) {
-        int i;
-
-        if (PyArray_NDIM(bmask) != 1) {
-            PyErr_SetString(PyExc_ValueError,
-                    "The boolean mask indexing array "
-                    "is neither one-dimensional nor "
-                    "matches the operand's number of "
-                    "dimensions");
-            return -1;
-        }
-        op_axes[1] = bmask_axes;
-        bmask_axes[0] = 0;
-        for (i = 1; i < PyArray_NDIM(self); ++i) {
-            bmask_axes[i] = -1;
-        }
+        PyErr_SetString(PyExc_ValueError,
+                "The boolean mask assignment indexing array "
+                "must have the same number of dimensions as "
+                "the array being indexed");
+        return -1;
     }
 
     /* See the Boolean Indexing section of the missing data NEP */
@@ -1086,9 +1051,8 @@ array_ass_boolean_subscript(PyArrayObject *self,
          */
         op_flags[1] = NPY_ITER_READONLY | NPY_ITER_IGNORE_MASKNA;
 
-        iter = NpyIter_AdvancedNew(2, op, flags, order, NPY_NO_CASTING,
-                                op_flags, NULL,
-                                PyArray_NDIM(self), op_axes, NULL, 0);
+        iter = NpyIter_MultiNew(2, op, flags, order, NPY_NO_CASTING,
+                                op_flags, NULL);
         if (iter == NULL) {
             return -1;
         }
@@ -1355,8 +1319,9 @@ array_subscript(PyArrayObject *self, PyObject *op)
         return NULL;
     }
 
-    /* Boolean indexing */
-    if (PyArray_Check(op) && (PyArray_TYPE((PyArrayObject *)op) == NPY_BOOL)) {
+    /* Boolean indexing special case which supports mask NA */
+    if (PyArray_Check(op) && (PyArray_TYPE((PyArrayObject *)op) == NPY_BOOL)
+                && (PyArray_NDIM(self) == PyArray_NDIM((PyArrayObject *)op))) {
         return (PyObject *)array_boolean_subscript(self,
                                         (PyArrayObject *)op, NPY_CORDER);
     }
@@ -1644,12 +1609,15 @@ array_ass_sub(PyArrayObject *self, PyObject *index, PyObject *op)
     }
     PyErr_Clear();
 
-    /* Boolean indexing */
+    /* Boolean indexing special case with NA mask support */
     if (PyArray_Check(index) &&
-                (PyArray_TYPE((PyArrayObject *)index) == NPY_BOOL)) {
+                (PyArray_TYPE((PyArrayObject *)index) == NPY_BOOL) &&
+                (PyArray_NDIM(self) == PyArray_NDIM((PyArrayObject *)index))) {
+        int retcode;
         PyArrayObject *op_arr;
         PyArray_Descr *dtype = NULL;
-        /* If it's an NA with no dtype, specify it explicitly */
+
+        /* If it's an NA with no dtype, specify the dtype explicitly */
         if (NpyNA_Check(op) && ((NpyNA_fields *)op)->dtype == NULL) {
             dtype = PyArray_DESCR(self);
             Py_INCREF(dtype);
@@ -1660,13 +1628,24 @@ array_ass_sub(PyArrayObject *self, PyObject *index, PyObject *op)
             return -1;
         }
 
-        return array_ass_boolean_subscript(self,
-                        (PyArrayObject *)index,
-                        (PyArrayObject *)op_arr, NPY_CORDER);
+        if (PyArray_NDIM(op_arr) < 2) {
+            retcode = array_ass_boolean_subscript(self,
+                            (PyArrayObject *)index,
+                            op_arr, NPY_CORDER);
+            Py_DECREF(op_arr);
+            return retcode;
+        }
+        /*
+         * Assigning from multi-dimensional 'op' in this case seems
+         * inconsistent, so falling through to old code for backwards
+         * compatibility.
+         */
+        Py_DECREF(op_arr);
     }
 
     fancy = fancy_indexing_check(index);
     if (fancy != SOBJ_NOTFANCY) {
+
         oned = ((PyArray_NDIM(self) == 1) &&
                 !(PyTuple_Check(index) && PyTuple_GET_SIZE(index) > 1));
         mit = (PyArrayMapIterObject *) PyArray_MapIterNew(index, oned, fancy);
@@ -1782,26 +1761,6 @@ array_subscript_nice(PyArrayObject *self, PyObject *op)
                                 PyArray_DESCR(self), (npy_mask)*maskna_item);
             }
         }
-
-#if 0
-        int i;
-        char *item;
-
-        for (i = 0; i < PyArray_NDIM(self); i++) {
-            if (vals[i] < 0) {
-                vals[i] += PyArray_DIMS(self)[i];
-            }
-            if ((vals[i] < 0) || (vals[i] >= PyArray_DIMS(self)[i])) {
-                PyErr_Format(PyExc_IndexError,
-                             "index (%"INTP_FMT") out of range "\
-                             "(0<=index<%"INTP_FMT") in dimension %d",
-                             vals[i], PyArray_DIMS(self)[i], i);
-                return NULL;
-            }
-        }
-        item = PyArray_GetPtr(self, vals);
-        return PyArray_Scalar(item, PyArray_DESCR(self), (PyObject *)self);
-#endif
     }
     PyErr_Clear();
 
@@ -1825,7 +1784,7 @@ array_subscript_nice(PyArrayObject *self, PyObject *op)
         return NULL;
     }
     if (PyArray_Check(mp) && PyArray_NDIM(mp) == 0) {
-        Bool noellipses = TRUE;
+        npy_bool noellipses = TRUE;
         if ((op == Py_Ellipsis) || PyString_Check(op) || PyUnicode_Check(op)) {
             noellipses = FALSE;
         }
@@ -2302,7 +2261,7 @@ PyArray_MapIterNew(PyObject *indexobj, int oned, int fancy)
         return NULL;
     }
 
-    mit = (PyArrayMapIterObject *)_pya_malloc(sizeof(PyArrayMapIterObject));
+    mit = (PyArrayMapIterObject *)PyArray_malloc(sizeof(PyArrayMapIterObject));
     PyObject_Init((PyObject *)mit, &PyArrayMapIter_Type);
     if (mit == NULL) {
         return NULL;
@@ -2362,7 +2321,7 @@ PyArray_MapIterNew(PyObject *indexobj, int oned, int fancy)
             PyTuple_SET_ITEM(mit->indexobj, i, PyInt_FromLong(0));
         }
     }
-    if (PyArray_Check(indexobj) || !PyTuple_Check(indexobj)) {
+    else if (PyArray_Check(indexobj) || !PyTuple_Check(indexobj)) {
         mit->numiter = 1;
         indtype = PyArray_DescrFromType(NPY_INTP);
         arr = (PyArrayObject *)PyArray_FromAny(indexobj, indtype, 0, 0,
@@ -2469,7 +2428,7 @@ arraymapiter_dealloc(PyArrayMapIterObject *mit)
     for (i = 0; i < mit->numiter; i++) {
         Py_XDECREF(mit->iters[i]);
     }
-    _pya_free(mit);
+    PyArray_free(mit);
 }
 
 /*
diff --git a/numpy/core/src/multiarray/shape.c b/numpy/core/src/multiarray/shape.c
index e2af79afd..9a3c29ad9 100644
--- a/numpy/core/src/multiarray/shape.c
+++ b/numpy/core/src/multiarray/shape.c
@@ -342,7 +342,7 @@ PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims,
                 }
             }
             else {
-                for (i = ndim; i >= 0; --i) {
+                for (i = ndim-1; i >= 0; --i) {
                     fa->maskna_strides[i] = stride;
                     stride *= fa->dimensions[i];
                 }
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index f7d775ca8..500c72620 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -4047,7 +4047,9 @@ PyUFunc_GenericReduction(PyUFuncObject *self, PyObject *args,
         if (axis < 0) {
             axis += PyArray_NDIM(mp);
         }
-        if (axis < 0 || axis >= PyArray_NDIM(mp)) {
+        /* Special case letting axis=0 slip through for scalars */
+        if ((axis < 0 || axis >= PyArray_NDIM(mp)) &&
+                !(axis == 0 && PyArray_NDIM(mp) == 0)) {
             PyErr_SetString(PyExc_ValueError,
                     "'axis' entry is out of bounds");
             Py_XDECREF(otype);
@@ -4060,8 +4062,14 @@ PyUFunc_GenericReduction(PyUFuncObject *self, PyObject *args,
 
     /* Check to see if input is zero-dimensional. */
     if (PyArray_NDIM(mp) == 0) {
-        /* A reduction with no axes is still valid but trivial */
-        if (operation == UFUNC_REDUCE && naxes == 0) {
+        /*
+         * A reduction with no axes is still valid but trivial.
+         * As a special case for backwards compatibility in 'sum',
+         * 'prod', et al, also allow a reduction where axis=0, even
+         * though this is technically incorrect.
+         */
+        if (operation == UFUNC_REDUCE &&
+                    (naxes == 0 || (naxes == 1 && axes[0] == 0))) {
             Py_XDECREF(otype);
             /* If there's an output parameter, copy the value */
             if (out != NULL) {
diff --git a/numpy/core/tests/test_indexing.py b/numpy/core/tests/test_indexing.py
new file mode 100644
index 000000000..da19d2a2a
--- /dev/null
+++ b/numpy/core/tests/test_indexing.py
@@ -0,0 +1,25 @@
+import numpy as np
+from numpy.compat import asbytes
+from numpy.testing import *
+import sys, warnings
+
+# The C implementation of fancy indexing is relatively complicated,
+# and has many seeming inconsistencies. It also appears to lack any
+# kind of test suite, making any changes to the underlying code difficult
+# because of its fragility.
+
+# This file is to remedy the test suite part a little bit,
+# but hopefully NumPy indexing can be changed to be more systematic
+# at some point in the future.
+
+def test_boolean_indexing():
+    # Indexing a 2-dimensional array with a length-1 array of 'True'
+    a = np.array([[ 0.,  0.,  0.]])
+    b = np.array([ True], dtype=bool)
+    assert_equal(a[b], a)
+
+    a[b] = 1.
+    assert_equal(a, [[1., 1., 1.]])
+
+if __name__ == "__main__":
+    run_module_suite()
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index ad483634d..69f2b3c74 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -505,6 +505,16 @@ class TestUfunc(TestCase):
         assert_raises(ValueError, np.max, [])
         assert_raises(ValueError, np.min, [])
 
+    def test_scalar_reduction(self):
+        # The functions 'sum', 'prod', etc allow specifying axis=0
+        # even for scalars
+        assert_equal(np.sum(3, axis=0), 3)
+        assert_equal(np.prod(3.5, axis=0), 3.5)
+        assert_equal(np.any(True, axis=0), True)
+        assert_equal(np.all(False, axis=0), False)
+        assert_equal(np.max(3, axis=0), 3)
+        assert_equal(np.min(2.5, axis=0), 2.5)
+
     def test_casting_out_param(self):
         # Test that it's possible to do casts on output
         a = np.ones((200,100), np.int64)
diff --git a/numpy/matrixlib/tests/test_defmatrix.py b/numpy/matrixlib/tests/test_defmatrix.py
index 09a4b4892..0a181fca3 100644
--- a/numpy/matrixlib/tests/test_defmatrix.py
+++ b/numpy/matrixlib/tests/test_defmatrix.py
@@ -65,29 +65,45 @@ class TestProperties(TestCase):
         sumall = 30
         assert_array_equal(sum0, M.sum(axis=0))
         assert_array_equal(sum1, M.sum(axis=1))
-        assert_(sumall == M.sum())
+        assert_equal(sumall, M.sum())
+
+        assert_array_equal(sum0, np.sum(M, axis=0))
+        assert_array_equal(sum1, np.sum(M, axis=1))
+        assert_equal(sumall, np.sum(M))
 
 
     def test_prod(self):
         x = matrix([[1,2,3],[4,5,6]])
-        assert_(x.prod() == 720)
-        assert_(all(x.prod(0) == matrix([[4,10,18]])))
-        assert_(all(x.prod(1) == matrix([[6],[120]])))
+        assert_equal(x.prod(), 720)
+        assert_equal(x.prod(0), matrix([[4,10,18]]))
+        assert_equal(x.prod(1), matrix([[6],[120]]))
+
+        assert_equal(np.prod(x), 720)
+        assert_equal(np.prod(x, axis=0), matrix([[4,10,18]]))
+        assert_equal(np.prod(x, axis=1), matrix([[6],[120]]))
 
         y = matrix([0,1,3])
         assert_(y.prod() == 0)
 
     def test_max(self):
         x = matrix([[1,2,3],[4,5,6]])
-        assert_(x.max() == 6)
-        assert_(all(x.max(0) == matrix([[4,5,6]])))
-        assert_(all(x.max(1) == matrix([[3],[6]])))
+        assert_equal(x.max(), 6)
+        assert_equal(x.max(0), matrix([[4,5,6]]))
+        assert_equal(x.max(1), matrix([[3],[6]]))
+
+        assert_equal(np.max(x), 6)
+        assert_equal(np.max(x, axis=0), matrix([[4,5,6]]))
+        assert_equal(np.max(x, axis=1), matrix([[3],[6]]))
 
     def test_min(self):
         x = matrix([[1,2,3],[4,5,6]])
-        assert_(x.min() == 1)
-        assert_(all(x.min(0) == matrix([[1,2,3]])))
-        assert_(all(x.min(1) == matrix([[1],[4]])))
+        assert_equal(x.min(), 1)
+        assert_equal(x.min(0), matrix([[1,2,3]]))
+        assert_equal(x.min(1), matrix([[1],[4]]))
+
+        assert_equal(np.min(x), 1)
+        assert_equal(np.min(x, axis=0), matrix([[1,2,3]]))
+        assert_equal(np.min(x, axis=1), matrix([[1],[4]]))
 
     def test_ptp(self):
         x = np.arange(4).reshape((2,2))