ENH: core: Add static caching of the callables for C to core._method calls

2 files changed, 147 insertions, 19 deletions

diff --git a/numpy/core/_methods.py b/numpy/core/_methods.py
new file mode 100644
index 000000000..3c05cf71d
--- /dev/null
+++ b/numpy/core/_methods.py
@@ -0,0 +1,105 @@
+# Array methods which are called by the both the C-code for the method
+# and the Python code for the NumPy-namespace function
+
+import multiarray as mu
+import umath as um
+from numeric import asanyarray
+
+def _amax(a, axis=None, out=None, skipna=False, keepdims=False):
+    return um.maximum.reduce(a, axis=axis,
+                            out=out, skipna=skipna, keepdims=keepdims)
+
+def _amin(a, axis=None, out=None, skipna=False, keepdims=False):
+    return um.minimum.reduce(a, axis=axis,
+                            out=out, skipna=skipna, keepdims=keepdims)
+
+def _sum(a, axis=None, dtype=None, out=None, skipna=False, keepdims=False):
+    return um.add.reduce(a, axis=axis, dtype=dtype,
+                            out=out, skipna=skipna, keepdims=keepdims)
+
+def _prod(a, axis=None, dtype=None, out=None, skipna=False, keepdims=False):
+    return um.multiply.reduce(a, axis=axis, dtype=dtype,
+                            out=out, skipna=skipna, keepdims=keepdims)
+
+def _any(a, axis=None, out=None, skipna=False, keepdims=False):
+    return um.logical_or.reduce(a, axis=axis,
+                            out=out, skipna=skipna, keepdims=keepdims)
+
+def _all(a, axis=None, out=None, skipna=False, keepdims=False):
+    return um.logical_and.reduce(a, axis=axis,
+                            out=out, skipna=skipna, keepdims=keepdims)
+
+def _mean(a, axis=None, dtype=None, out=None, skipna=False, keepdims=False):
+    arr = asanyarray(a)
+
+    # Upgrade bool, unsigned int, and int to float64
+    if dtype is None and arr.dtype.kind in ['b','u','i']:
+        ret = um.add.reduce(arr, axis=axis, dtype='f8',
+                            out=out, skipna=skipna, keepdims=keepdims)
+    else:
+        ret = um.add.reduce(arr, axis=axis, dtype=dtype,
+                            out=out, skipna=skipna, keepdims=keepdims)
+    rcount = mu.count_reduce_items(arr, axis=axis,
+                            skipna=skipna, keepdims=keepdims)
+    if isinstance(ret, mu.ndarray):
+        ret = um.true_divide(ret, rcount,
+                        out=ret, casting='unsafe', subok=False)
+    else:
+        ret = ret / float(rcount)
+    return ret
+
+def _var(a, axis=None, dtype=None, out=None, ddof=0,
+                            skipna=False, keepdims=False):
+    arr = asanyarray(a)
+
+    # First compute the mean, saving 'rcount' for reuse later
+    if dtype is None and arr.dtype.kind in ['b','u','i']:
+        arrmean = um.add.reduce(arr, axis=axis, dtype='f8',
+                            skipna=skipna, keepdims=True)
+    else:
+        arrmean = um.add.reduce(arr, axis=axis, dtype=dtype,
+                            skipna=skipna, keepdims=True)
+    rcount = mu.count_reduce_items(arr, axis=axis,
+                            skipna=skipna, keepdims=True)
+    if isinstance(arrmean, mu.ndarray):
+        arrmean = um.true_divide(arrmean, rcount,
+                            out=arrmean, casting='unsafe', subok=False)
+    else:
+        arrmean = arrmean / float(rcount)
+
+    # arr - arrmean
+    x = arr - arrmean
+
+    # (arr - arrmean) ** 2
+    if arr.dtype.kind == 'c':
+        x = um.multiply(x, um.conjugate(x), out=x).real
+    else:
+        x = um.multiply(x, x, out=x)
+
+    # add.reduce((arr - arrmean) ** 2, axis)
+    ret = um.add.reduce(x, axis=axis, dtype=dtype, out=out,
+                                skipna=skipna, keepdims=keepdims)
+
+    # add.reduce((arr - arrmean) ** 2, axis) / (n - ddof)
+    if not keepdims and isinstance(rcount, mu.ndarray):
+        rcount = rcount.squeeze(axis=axis)
+    rcount -= ddof
+    if isinstance(ret, mu.ndarray):
+        ret = um.true_divide(ret, rcount,
+                        out=ret, casting='unsafe', subok=False)
+    else:
+        ret = ret / float(rcount)
+
+    return ret
+
+def _std(a, axis=None, dtype=None, out=None, ddof=0,
+                            skipna=False, keepdims=False):
+    ret = _var(a, axis=axis, dtype=dtype, out=out, ddof=ddof,
+                                skipna=skipna, keepdims=keepdims)
+
+    if isinstance(ret, mu.ndarray):
+        um.sqrt(ret, out=ret)
+    else:
+        ret = um.sqrt(ret)
+
+    return ret
diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
index afeae9a30..6af1ca984 100644
--- a/numpy/core/src/multiarray/methods.c
+++ b/numpy/core/src/multiarray/methods.c
@@ -49,17 +49,10 @@ NpyArg_ParseKeywords(PyObject *keys, const char *format, char **kwlist, ...)
     return ret;
 }
 
-/*
- * Forwards an ndarray method to a the Python function
- * numpy.core._methods.<name>(...)
- */
 static PyObject *
-forward_ndarray_method(PyArrayObject *self, PyObject *args, PyObject *kwds,
-                            const char *name)
+get_forwarding_ndarray_method(const char *name)
 {
-    PyObject *sargs, *ret;
     PyObject *module_methods, *callable;
-    int i, n;
 
     /* Get a reference to the function we're calling */
     module_methods = PyImport_ImportModule("numpy.core._methods");
@@ -72,14 +65,28 @@ forward_ndarray_method(PyArrayObject *self, PyObject *args, PyObject *kwds,
         PyErr_Format(PyExc_RuntimeError,
                 "NumPy internal error: could not find function "
                 "numpy.core._methods.%s", name);
-        return NULL;
     }
 
+    Py_INCREF(callable);
+    Py_DECREF(module_methods);
+    return callable;
+}
+
+/*
+ * Forwards an ndarray method to a the Python function
+ * numpy.core._methods.<name>(...)
+ */
+static PyObject *
+forward_ndarray_method(PyArrayObject *self, PyObject *args, PyObject *kwds,
+                            PyObject *forwarding_callable)
+{
+    PyObject *sargs, *ret;
+    int i, n;
+
     /* Combine 'self' and 'args' together into one tuple */
     n = PyTuple_GET_SIZE(args);
     sargs = PyTuple_New(n + 1);
     if (sargs == NULL) {
-        Py_DECREF(module_methods);
         return NULL;
     }
     Py_INCREF(self);
@@ -91,12 +98,28 @@ forward_ndarray_method(PyArrayObject *self, PyObject *args, PyObject *kwds,
     }
 
     /* Call the function and return */
-    ret = PyObject_Call(callable, sargs, kwds);
+    ret = PyObject_Call(forwarding_callable, sargs, kwds);
     Py_DECREF(sargs);
-    Py_DECREF(module_methods);
     return ret;
 }
 
+/*
+ * Forwards an ndarray method to the function numpy.core._methods.<name>(...),
+ * caching the callable in a local static variable. Note that the
+ * initialization is not thread-safe, but relies on the CPython GIL to
+ * be correct.
+ */
+#define NPY_FORWARD_NDARRAY_METHOD(name) \
+        static PyObject *callable = NULL; \
+        if (callable == NULL) { \
+            callable = get_forwarding_ndarray_method(name); \
+            if (callable == NULL) { \
+                return NULL; \
+            } \
+        } \
+        return forward_ndarray_method(self, args, kwds, callable)
+
+
 static PyObject *
 array_take(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
@@ -340,13 +363,13 @@ array_argmin(PyArrayObject *self, PyObject *args, PyObject *kwds)
 static PyObject *
 array_max(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
-    return forward_ndarray_method(self, args, kwds, "_amax");
+    NPY_FORWARD_NDARRAY_METHOD("_amax");
 }
 
 static PyObject *
 array_min(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
-    return forward_ndarray_method(self, args, kwds, "_amin");
+    NPY_FORWARD_NDARRAY_METHOD("_amin");
 }
 
 static PyObject *
@@ -1885,13 +1908,13 @@ _get_type_num_double(PyArray_Descr *dtype1, PyArray_Descr *dtype2)
 static PyObject *
 array_mean(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
-    return forward_ndarray_method(self, args, kwds, "_mean");
+    NPY_FORWARD_NDARRAY_METHOD("_mean");
 }
 
 static PyObject *
 array_sum(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
-    return forward_ndarray_method(self, args, kwds, "_sum");
+    NPY_FORWARD_NDARRAY_METHOD("_sum");
 }
 
 
@@ -1920,7 +1943,7 @@ array_cumsum(PyArrayObject *self, PyObject *args, PyObject *kwds)
 static PyObject *
 array_prod(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
-    return forward_ndarray_method(self, args, kwds, "_prod");
+    NPY_FORWARD_NDARRAY_METHOD("_prod");
 }
 
 static PyObject *
@@ -2004,14 +2027,14 @@ array_all(PyArrayObject *self, PyObject *args, PyObject *kwds)
 static PyObject *
 array_stddev(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
-    return forward_ndarray_method(self, args, kwds, "_std");
+    NPY_FORWARD_NDARRAY_METHOD("_std");
 }
 
 
 static PyObject *
 array_variance(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
-    return forward_ndarray_method(self, args, kwds, "_var");
+    NPY_FORWARD_NDARRAY_METHOD("_var");
 }