summaryrefslogtreecommitdiff
path: root/numpy/core
diff options
context:
space:
mode:
Diffstat (limited to 'numpy/core')
-rw-r--r--numpy/core/src/multiarray/calculation.c190
-rw-r--r--numpy/core/tests/test_maskna.py17
2 files changed, 134 insertions, 73 deletions
diff --git a/numpy/core/src/multiarray/calculation.c b/numpy/core/src/multiarray/calculation.c
index a6cb82ecb..9ddbc4b89 100644
--- a/numpy/core/src/multiarray/calculation.c
+++ b/numpy/core/src/multiarray/calculation.c
@@ -183,7 +183,8 @@ PyArray_Max(PyArrayObject *ap, int axis, PyArrayObject *out)
PyArrayObject *arr;
PyObject *ret;
- if ((arr=(PyArrayObject *)PyArray_CheckAxis(ap, &axis, 0)) == NULL) {
+ arr = (PyArrayObject *)PyArray_CheckAxis(ap, &axis, NPY_ARRAY_ALLOWNA);
+ if (arr == NULL) {
return NULL;
}
ret = PyArray_GenericReduceFunction(arr, n_ops.maximum, axis,
@@ -201,7 +202,8 @@ PyArray_Min(PyArrayObject *ap, int axis, PyArrayObject *out)
PyArrayObject *arr;
PyObject *ret;
- if ((arr=(PyArrayObject *)PyArray_CheckAxis(ap, &axis, 0)) == NULL) {
+ arr=(PyArrayObject *)PyArray_CheckAxis(ap, &axis, NPY_ARRAY_ALLOWNA);
+ if (arr == NULL) {
return NULL;
}
ret = PyArray_GenericReduceFunction(arr, n_ops.minimum, axis,
@@ -220,7 +222,8 @@ PyArray_Ptp(PyArrayObject *ap, int axis, PyArrayObject *out)
PyObject *ret;
PyObject *obj1 = NULL, *obj2 = NULL;
- if ((arr=(PyArrayObject *)PyArray_CheckAxis(ap, &axis, 0)) == NULL) {
+ arr=(PyArrayObject *)PyArray_CheckAxis(ap, &axis, NPY_ARRAY_ALLOWNA);
+ if (arr == NULL) {
return NULL;
}
obj1 = PyArray_Max(arr, axis, out);
@@ -267,26 +270,27 @@ __New_PyArray_Std(PyArrayObject *self, int axis, int rtype, PyArrayObject *out,
int variance, int num)
{
PyObject *obj1 = NULL, *obj2 = NULL, *obj3 = NULL;
- PyArrayObject *arr1 = NULL, *arr2 = NULL, *new = NULL;
+ PyArrayObject *arr1 = NULL, *arr2 = NULL, *arrnew = NULL;
PyObject *ret = NULL, *newshape = NULL;
int i, n;
intp val;
- if ((new = (PyArrayObject *)PyArray_CheckAxis(self, &axis, 0)) == NULL) {
+ arrnew = (PyArrayObject *)PyArray_CheckAxis(self, &axis, NPY_ARRAY_ALLOWNA);
+ if (arrnew == NULL) {
return NULL;
}
/* Compute and reshape mean */
arr1 = (PyArrayObject *)PyArray_EnsureAnyArray(
- PyArray_Mean(new, axis, rtype, NULL));
+ PyArray_Mean(arrnew, axis, rtype, NULL));
if (arr1 == NULL) {
- Py_DECREF(new);
+ Py_DECREF(arrnew);
return NULL;
}
- n = PyArray_NDIM(new);
+ n = PyArray_NDIM(arrnew);
newshape = PyTuple_New(n);
if (newshape == NULL) {
Py_DECREF(arr1);
- Py_DECREF(new);
+ Py_DECREF(arrnew);
return NULL;
}
for (i = 0; i < n; i++) {
@@ -294,7 +298,7 @@ __New_PyArray_Std(PyArrayObject *self, int axis, int rtype, PyArrayObject *out,
val = 1;
}
else {
- val = PyArray_DIM(new,i);
+ val = PyArray_DIM(arrnew,i);
}
PyTuple_SET_ITEM(newshape, i, PyInt_FromLong((long)val));
}
@@ -302,16 +306,16 @@ __New_PyArray_Std(PyArrayObject *self, int axis, int rtype, PyArrayObject *out,
Py_DECREF(arr1);
Py_DECREF(newshape);
if (arr2 == NULL) {
- Py_DECREF(new);
+ Py_DECREF(arrnew);
return NULL;
}
/* Compute x = x - mx */
arr1 = (PyArrayObject *)PyArray_EnsureAnyArray(
- PyNumber_Subtract((PyObject *)new, (PyObject *)arr2));
+ PyNumber_Subtract((PyObject *)arrnew, (PyObject *)arr2));
Py_DECREF(arr2);
if (arr1 == NULL) {
- Py_DECREF(new);
+ Py_DECREF(arrnew);
return NULL;
}
/* Compute x * x */
@@ -323,7 +327,7 @@ __New_PyArray_Std(PyArrayObject *self, int axis, int rtype, PyArrayObject *out,
Py_INCREF(arr1);
}
if (obj3 == NULL) {
- Py_DECREF(new);
+ Py_DECREF(arrnew);
return NULL;
}
arr2 = (PyArrayObject *)PyArray_EnsureAnyArray(
@@ -331,7 +335,7 @@ __New_PyArray_Std(PyArrayObject *self, int axis, int rtype, PyArrayObject *out,
Py_DECREF(arr1);
Py_DECREF(obj3);
if (arr2 == NULL) {
- Py_DECREF(new);
+ Py_DECREF(arrnew);
return NULL;
}
if (PyArray_ISCOMPLEX(arr2)) {
@@ -353,7 +357,7 @@ __New_PyArray_Std(PyArrayObject *self, int axis, int rtype, PyArrayObject *out,
Py_INCREF(arr2);
}
if (obj3 == NULL) {
- Py_DECREF(new);
+ Py_DECREF(arrnew);
return NULL;
}
/* Compute add.reduce(x*x,axis) */
@@ -362,11 +366,11 @@ __New_PyArray_Std(PyArrayObject *self, int axis, int rtype, PyArrayObject *out,
Py_DECREF(obj3);
Py_DECREF(arr2);
if (obj1 == NULL) {
- Py_DECREF(new);
+ Py_DECREF(arrnew);
return NULL;
}
- n = PyArray_DIM(new,axis);
- Py_DECREF(new);
+ n = PyArray_DIM(arrnew,axis);
+ Py_DECREF(arrnew);
n = (n-num);
if (n == 0) {
n = 1;
@@ -422,14 +426,15 @@ finish:
NPY_NO_EXPORT PyObject *
PyArray_Sum(PyArrayObject *self, int axis, int rtype, PyArrayObject *out)
{
- PyObject *new, *ret;
+ PyObject *arr, *ret;
- if ((new = PyArray_CheckAxis(self, &axis, 0)) == NULL) {
+ arr = PyArray_CheckAxis(self, &axis, NPY_ARRAY_ALLOWNA);
+ if (arr == NULL) {
return NULL;
}
- ret = PyArray_GenericReduceFunction((PyArrayObject *)new, n_ops.add, axis,
+ ret = PyArray_GenericReduceFunction((PyArrayObject *)arr, n_ops.add, axis,
rtype, out);
- Py_DECREF(new);
+ Py_DECREF(arr);
return ret;
}
@@ -439,14 +444,16 @@ PyArray_Sum(PyArrayObject *self, int axis, int rtype, PyArrayObject *out)
NPY_NO_EXPORT PyObject *
PyArray_Prod(PyArrayObject *self, int axis, int rtype, PyArrayObject *out)
{
- PyObject *new, *ret;
+ PyObject *arr, *ret;
- if ((new = PyArray_CheckAxis(self, &axis, 0)) == NULL) {
+ arr = PyArray_CheckAxis(self, &axis, NPY_ARRAY_ALLOWNA);
+ if (arr == NULL) {
return NULL;
}
- ret = PyArray_GenericReduceFunction((PyArrayObject *)new, n_ops.multiply, axis,
+ ret = PyArray_GenericReduceFunction((PyArrayObject *)arr,
+ n_ops.multiply, axis,
rtype, out);
- Py_DECREF(new);
+ Py_DECREF(arr);
return ret;
}
@@ -456,14 +463,16 @@ PyArray_Prod(PyArrayObject *self, int axis, int rtype, PyArrayObject *out)
NPY_NO_EXPORT PyObject *
PyArray_CumSum(PyArrayObject *self, int axis, int rtype, PyArrayObject *out)
{
- PyObject *new, *ret;
+ PyObject *arr, *ret;
- if ((new = PyArray_CheckAxis(self, &axis, 0)) == NULL) {
+ arr = PyArray_CheckAxis(self, &axis, NPY_ARRAY_ALLOWNA);
+ if (arr == NULL) {
return NULL;
}
- ret = PyArray_GenericAccumulateFunction((PyArrayObject *)new, n_ops.add, axis,
+ ret = PyArray_GenericAccumulateFunction((PyArrayObject *)arr,
+ n_ops.add, axis,
rtype, out);
- Py_DECREF(new);
+ Py_DECREF(arr);
return ret;
}
@@ -473,16 +482,17 @@ PyArray_CumSum(PyArrayObject *self, int axis, int rtype, PyArrayObject *out)
NPY_NO_EXPORT PyObject *
PyArray_CumProd(PyArrayObject *self, int axis, int rtype, PyArrayObject *out)
{
- PyObject *new, *ret;
+ PyObject *arr, *ret;
- if ((new = PyArray_CheckAxis(self, &axis, 0)) == NULL) {
+ arr = PyArray_CheckAxis(self, &axis, NPY_ARRAY_ALLOWNA);
+ if (arr == NULL) {
return NULL;
}
- ret = PyArray_GenericAccumulateFunction((PyArrayObject *)new,
+ ret = PyArray_GenericAccumulateFunction((PyArrayObject *)arr,
n_ops.multiply, axis,
rtype, out);
- Py_DECREF(new);
+ Py_DECREF(arr);
return ret;
}
@@ -503,24 +513,24 @@ PyArray_Round(PyArrayObject *a, int decimals, PyArrayObject *out)
if (PyArray_ISCOMPLEX(a)) {
PyObject *part;
PyObject *round_part;
- PyObject *new;
+ PyObject *arr;
int res;
if (out) {
- new = (PyObject *)out;
- Py_INCREF(new);
+ arr = (PyObject *)out;
+ Py_INCREF(arr);
}
else {
- new = PyArray_Copy(a);
- if (new == NULL) {
+ arr = PyArray_Copy(a);
+ if (arr == NULL) {
return NULL;
}
}
- /* new.real = a.real.round(decimals) */
- part = PyObject_GetAttrString(new, "real");
+ /* arr.real = a.real.round(decimals) */
+ part = PyObject_GetAttrString(arr, "real");
if (part == NULL) {
- Py_DECREF(new);
+ Py_DECREF(arr);
return NULL;
}
part = PyArray_EnsureAnyArray(part);
@@ -528,20 +538,20 @@ PyArray_Round(PyArrayObject *a, int decimals, PyArrayObject *out)
decimals, NULL);
Py_DECREF(part);
if (round_part == NULL) {
- Py_DECREF(new);
+ Py_DECREF(arr);
return NULL;
}
- res = PyObject_SetAttrString(new, "real", round_part);
+ res = PyObject_SetAttrString(arr, "real", round_part);
Py_DECREF(round_part);
if (res < 0) {
- Py_DECREF(new);
+ Py_DECREF(arr);
return NULL;
}
- /* new.imag = a.imag.round(decimals) */
- part = PyObject_GetAttrString(new, "imag");
+ /* arr.imag = a.imag.round(decimals) */
+ part = PyObject_GetAttrString(arr, "imag");
if (part == NULL) {
- Py_DECREF(new);
+ Py_DECREF(arr);
return NULL;
}
part = PyArray_EnsureAnyArray(part);
@@ -549,16 +559,16 @@ PyArray_Round(PyArrayObject *a, int decimals, PyArrayObject *out)
decimals, NULL);
Py_DECREF(part);
if (round_part == NULL) {
- Py_DECREF(new);
+ Py_DECREF(arr);
return NULL;
}
- res = PyObject_SetAttrString(new, "imag", round_part);
+ res = PyObject_SetAttrString(arr, "imag", round_part);
Py_DECREF(round_part);
if (res < 0) {
- Py_DECREF(new);
+ Py_DECREF(arr);
return NULL;
}
- return new;
+ return arr;
}
/* do the most common case first */
if (decimals >= 0) {
@@ -652,15 +662,16 @@ NPY_NO_EXPORT PyObject *
PyArray_Mean(PyArrayObject *self, int axis, int rtype, PyArrayObject *out)
{
PyObject *obj1 = NULL, *obj2 = NULL, *ret;
- PyArrayObject *new;
+ PyArrayObject *arr;
- if ((new = (PyArrayObject *)PyArray_CheckAxis(self, &axis, 0)) == NULL) {
+ arr = (PyArrayObject *)PyArray_CheckAxis(self, &axis, NPY_ARRAY_ALLOWNA);
+ if (arr == NULL) {
return NULL;
}
- obj1 = PyArray_GenericReduceFunction(new, n_ops.add, axis,
+ obj1 = PyArray_GenericReduceFunction(arr, n_ops.add, axis,
rtype, out);
- obj2 = PyFloat_FromDouble((double) PyArray_DIM(new,axis));
- Py_DECREF(new);
+ obj2 = PyFloat_FromDouble((double)PyArray_DIM(arr,axis));
+ Py_DECREF(arr);
if (obj1 == NULL || obj2 == NULL) {
Py_XDECREF(obj1);
Py_XDECREF(obj2);
@@ -687,15 +698,16 @@ PyArray_Mean(PyArrayObject *self, int axis, int rtype, PyArrayObject *out)
NPY_NO_EXPORT PyObject *
PyArray_Any(PyArrayObject *self, int axis, PyArrayObject *out)
{
- PyObject *new, *ret;
+ PyObject *arr, *ret;
- if ((new = PyArray_CheckAxis(self, &axis, 0)) == NULL) {
+ arr = PyArray_CheckAxis(self, &axis, NPY_ARRAY_ALLOWNA);
+ if (arr == NULL) {
return NULL;
}
- ret = PyArray_GenericReduceFunction((PyArrayObject *)new,
+ ret = PyArray_GenericReduceFunction((PyArrayObject *)arr,
n_ops.logical_or, axis,
PyArray_BOOL, out);
- Py_DECREF(new);
+ Py_DECREF(arr);
return ret;
}
@@ -705,15 +717,16 @@ PyArray_Any(PyArrayObject *self, int axis, PyArrayObject *out)
NPY_NO_EXPORT PyObject *
PyArray_All(PyArrayObject *self, int axis, PyArrayObject *out)
{
- PyObject *new, *ret;
+ PyObject *arr, *ret;
- if ((new = PyArray_CheckAxis(self, &axis, 0)) == NULL) {
+ arr = PyArray_CheckAxis(self, &axis, NPY_ARRAY_ALLOWNA);
+ if (arr == NULL) {
return NULL;
}
- ret = PyArray_GenericReduceFunction((PyArrayObject *)new,
+ ret = PyArray_GenericReduceFunction((PyArrayObject *)arr,
n_ops.logical_and, axis,
PyArray_BOOL, out);
- Py_DECREF(new);
+ Py_DECREF(arr);
return ret;
}
@@ -787,6 +800,11 @@ _slow_array_clip(PyArrayObject *self, PyObject *min, PyObject *max, PyArrayObjec
/*NUMPY_API
* Clip
+ *
+ * TODO: For adding NA support, a Clip UFunc should be created, then
+ * this should call that ufunc. 'min' and 'max' can default to
+ * the -inf/+inf or the smallest/largest representable values
+ * of the dtype respectively.
*/
NPY_NO_EXPORT PyObject *
PyArray_Clip(PyArrayObject *self, PyObject *min, PyObject *max, PyArrayObject *out)
@@ -800,15 +818,32 @@ PyArray_Clip(PyArrayObject *self, PyObject *min, PyObject *max, PyArrayObject *o
char *max_data, *min_data;
PyObject *zero;
+ /* Treat None the same as NULL */
+ if (min == Py_None) {
+ min = NULL;
+ }
+ if (max == Py_None) {
+ max = NULL;
+ }
+
if ((max == NULL) && (min == NULL)) {
- PyErr_SetString(PyExc_ValueError, "array_clip: must set either max "\
- "or min");
+ PyErr_SetString(PyExc_ValueError,
+ "array_clip: must set either max or min");
return NULL;
}
func = PyArray_DESCR(self)->f->fastclip;
- if (func == NULL || (min != NULL && !PyArray_CheckAnyScalar(min)) ||
- (max != NULL && !PyArray_CheckAnyScalar(max))) {
+ /* Trigger the slow array clip for NA support as well */
+ if (func == NULL ||
+ PyArray_HASMASKNA(self) ||
+ (min != NULL &&
+ (!PyArray_CheckAnyScalar(min) ||
+ (PyArray_Check(min) &&
+ PyArray_HASMASKNA((PyArrayObject *)min)))) ||
+ (max != NULL &&
+ (!PyArray_CheckAnyScalar(max) ||
+ (PyArray_Check(max) &&
+ PyArray_HASMASKNA((PyArrayObject *)max))))) {
return _slow_array_clip(self, min, max, out);
}
/* Use the fast scalar clip function */
@@ -866,7 +901,7 @@ PyArray_Clip(PyArrayObject *self, PyObject *min, PyObject *max, PyArrayObject *o
/* Convert max to an array */
if (max != NULL) {
maxa = (PyArrayObject *)PyArray_FromAny(max, indescr, 0, 0,
- NPY_ARRAY_DEFAULT, NULL);
+ NPY_ARRAY_DEFAULT, NULL);
if (maxa == NULL) {
return NULL;
}
@@ -908,7 +943,7 @@ PyArray_Clip(PyArrayObject *self, PyObject *min, PyObject *max, PyArrayObject *o
/* Convert min to an array */
Py_INCREF(indescr);
mina = (PyArrayObject *)PyArray_FromAny(min, indescr, 0, 0,
- NPY_ARRAY_DEFAULT, NULL);
+ NPY_ARRAY_DEFAULT, NULL);
Py_DECREF(min);
if (mina == NULL) {
goto fail;
@@ -978,6 +1013,13 @@ PyArray_Clip(PyArrayObject *self, PyObject *min, PyObject *max, PyArrayObject *o
if (out == NULL) {
goto fail;
}
+
+ if ((maxa != NULL && PyArray_HASMASKNA(maxa)) ||
+ (mina != NULL && PyArray_HASMASKNA(mina))) {
+ if (PyArray_AllocateMaskNA(out, 1, 0, 1) < 0) {
+ goto fail;
+ }
+ }
outgood = 1;
}
else Py_INCREF(out);
@@ -1021,7 +1063,9 @@ PyArray_Clip(PyArrayObject *self, PyObject *min, PyObject *max, PyArrayObject *o
goto fail;
}
if (PyArray_DATA(newout) != PyArray_DATA(newin)) {
- memcpy(PyArray_DATA(newout), PyArray_DATA(newin), PyArray_NBYTES(newin));
+ if (PyArray_CopyInto(newout, newin) < 0) {
+ goto fail;
+ }
}
/* Now we can call the fast-clip function */
diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py
index e5d8d0d14..1a5561c3b 100644
--- a/numpy/core/tests/test_maskna.py
+++ b/numpy/core/tests/test_maskna.py
@@ -719,5 +719,22 @@ def check_ufunc_max_1D(max_func):
a[...] = np.NA
assert_raises(ValueError, max_func, a, skipna=True)
+def test_array_maskna_methods():
+ a = np.array([2, np.NA, 10, 4, np.NA, 7], maskna=True)
+
+ # ndarray.clip
+ b = np.clip(a, 3, None)
+ print repr(b)
+ assert_equal(np.isna(b), [0,1,0,0,1,0])
+ assert_equal(b[~np.isna(b)], [3, 10, 4, 7])
+
+ b = np.clip(a, None, 6)
+ assert_equal(np.isna(b), [0,1,0,0,1,0])
+ assert_equal(b[~np.isna(b)], [2, 6, 4, 6])
+
+ b = np.clip(a, 4, 7)
+ assert_equal(np.isna(b), [0,1,0,0,1,0])
+ assert_equal(b[~np.isna(b)], [4, 7, 4, 7])
+
if __name__ == "__main__":
run_module_suite()
View Lorry Controller Status