numpy · ahaldane · Sep 18, 2017 · Jun 2, 2017 · Jun 2, 2017 · Jul 1, 2017
diff --git a/doc/release/1.14.0-notes.rst b/doc/release/1.14.0-notes.rst
@@ -217,6 +217,11 @@ selected via the ``--fcompiler`` and ``--compiler`` options to
 supported; by default a gfortran-compatible static archive
 ``openblas.a`` is looked for.
 
+``concatenate`` and ``stack`` gained an ``out`` argument
+--------------------------------------------------------
+A preallocated buffer of the desired dtype can now be used for the output of
+these functions.
+
 ``np.linalg.pinv`` now works on stacked matrices
 ------------------------------------------------
 Previously it was limited to a single 2d array.

diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
@@ -1158,7 +1158,7 @@ def luf(lamdaexpr, *args, **kwargs):
 
 add_newdoc('numpy.core.multiarray', 'concatenate',
     """
-    concatenate((a1, a2, ...), axis=0)
+    concatenate((a1, a2, ...), axis=0, out=None)
 
     Join a sequence of arrays along an existing axis.
 
@@ -1169,6 +1169,10 @@ def luf(lamdaexpr, *args, **kwargs):
         corresponding to `axis` (the first, by default).
     axis : int, optional
         The axis along which the arrays will be joined.  Default is 0.
+    out : ndarray, optional
+        If provided, the destination to place the result. The shape must be
+        correct, matching that of what concatenate would have returned if no
+        out argument were specified.
 
     Returns
     -------

diff --git a/numpy/core/shape_base.py b/numpy/core/shape_base.py
@@ -293,7 +293,7 @@ def hstack(tup):
         return _nx.concatenate(arrs, 1)
 
 
-def stack(arrays, axis=0):
+def stack(arrays, axis=0, out=None):
     """
     Join a sequence of arrays along a new axis.
 
@@ -309,6 +309,10 @@ def stack(arrays, axis=0):
         Each array must have the same shape.
     axis : int, optional
         The axis in the result array along which the input arrays are stacked.
+    out : ndarray, optional
+        If provided, the destination to place the result. The shape must be
+        correct, matching that of what stack would have returned if no
+        out argument were specified.
 
     Returns
     -------
@@ -358,7 +362,7 @@ def stack(arrays, axis=0):
 
     sl = (slice(None),) * axis + (_nx.newaxis,)
     expanded_arrays = [arr[sl] for arr in arrays]
-    return _nx.concatenate(expanded_arrays, axis=axis)
+    return _nx.concatenate(expanded_arrays, axis=axis, out=out)
 
 
 class _Recurser(object):

diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
@@ -315,20 +315,39 @@ PyArray_Free(PyObject *op, void *ptr)
     return 0;
 }
 
+/*
+ * Get the ndarray subclass with the highest priority
+ */
+NPY_NO_EXPORT PyTypeObject *
+PyArray_GetSubType(int narrays, PyArrayObject **arrays) {
+    PyTypeObject *subtype = &PyArray_Type;
+    double priority = NPY_PRIORITY;
+    int i;
+
+    /* Get the priority subtype for the array */
+    for (i = 0; i < narrays; ++i) {
+        if (Py_TYPE(arrays[i]) != subtype) {
+            double pr = PyArray_GetPriority((PyObject *)(arrays[i]), 0.0);
+            if (pr > priority) {
+                priority = pr;
+                subtype = Py_TYPE(arrays[i]);
+            }
+        }
+    }
+
+    return subtype;
+}
+
 
 /*
  * Concatenates a list of ndarrays.
  */
 NPY_NO_EXPORT PyArrayObject *
-PyArray_ConcatenateArrays(int narrays, PyArrayObject **arrays, int axis)
+PyArray_ConcatenateArrays(int narrays, PyArrayObject **arrays, int axis,
+                          PyArrayObject* ret)
 {
-    PyTypeObject *subtype = &PyArray_Type;
-    double priority = NPY_PRIORITY;
     int iarrays, idim, ndim;
-    npy_intp shape[NPY_MAXDIMS], s, strides[NPY_MAXDIMS];
-    int strideperm[NPY_MAXDIMS];
-    PyArray_Descr *dtype = NULL;
-    PyArrayObject *ret = NULL;
+    npy_intp shape[NPY_MAXDIMS];
     PyArrayObject_fields *sliding_view = NULL;
 
     if (narrays <= 0) {
@@ -383,47 +402,57 @@ PyArray_ConcatenateArrays(int narrays, PyArrayObject **arrays, int axis)
         }
     }
 
-    /* Get the priority subtype for the array */
-    for (iarrays = 0; iarrays < narrays; ++iarrays) {
-        if (Py_TYPE(arrays[iarrays]) != subtype) {
-            double pr = PyArray_GetPriority((PyObject *)(arrays[iarrays]), 0.0);
-            if (pr > priority) {
-                priority = pr;
-                subtype = Py_TYPE(arrays[iarrays]);
-            }
+    if (ret != NULL) {
+        if (PyArray_NDIM(ret) != ndim) {
+            PyErr_SetString(PyExc_ValueError,
+                            "Output array has wrong dimensionality");
+            return NULL;
         }
+        if (!PyArray_CompareLists(shape, PyArray_SHAPE(ret), ndim)) {
+            PyErr_SetString(PyExc_ValueError,
+                            "Output array is the wrong shape");
+            return NULL;
+        }
+        Py_INCREF(ret);
     }
+    else {
+        npy_intp s, strides[NPY_MAXDIMS];
+        int strideperm[NPY_MAXDIMS];
 
-    /* Get the resulting dtype from combining all the arrays */
-    dtype = PyArray_ResultType(narrays, arrays, 0, NULL);
-    if (dtype == NULL) {
-        return NULL;
-    }
+        /* Get the priority subtype for the array */
+        PyTypeObject *subtype = PyArray_GetSubType(narrays, arrays);
 
-    /*
-     * Figure out the permutation to apply to the strides to match
-     * the memory layout of the input arrays, using ambiguity
-     * resolution rules matching that of the NpyIter.
-     */
-    PyArray_CreateMultiSortedStridePerm(narrays, arrays, ndim, strideperm);
-    s = dtype->elsize;
-    for (idim = ndim-1; idim >= 0; --idim) {
-        int iperm = strideperm[idim];
-        strides[iperm] = s;
-        s *= shape[iperm];
-    }
-
-    /* Allocate the array for the result. This steals the 'dtype' reference. */
-    ret = (PyArrayObject *)PyArray_NewFromDescr(subtype,
-                                                    dtype,
-                                                    ndim,
-                                                    shape,
-                                                    strides,
-                                                    NULL,
-                                                    0,
-                                                    NULL);
-    if (ret == NULL) {
-        return NULL;
+        /* Get the resulting dtype from combining all the arrays */
+        PyArray_Descr *dtype = PyArray_ResultType(narrays, arrays, 0, NULL);
+        if (dtype == NULL) {
+            return NULL;
+        }
+
+        /*
+         * Figure out the permutation to apply to the strides to match
+         * the memory layout of the input arrays, using ambiguity
+         * resolution rules matching that of the NpyIter.
+         */
+        PyArray_CreateMultiSortedStridePerm(narrays, arrays, ndim, strideperm);
+        s = dtype->elsize;
+        for (idim = ndim-1; idim >= 0; --idim) {
+            int iperm = strideperm[idim];
+            strides[iperm] = s;
+            s *= shape[iperm];
+        }
+
+        /* Allocate the array for the result. This steals the 'dtype' reference. */
+        ret = (PyArrayObject *)PyArray_NewFromDescr(subtype,
+                                                        dtype,
+                                                        ndim,
+                                                        shape,
+                                                        strides,
+                                                        NULL,
+                                                        0,
+                                                        NULL);
+        if (ret == NULL) {
+            return NULL;
+        }
     }
 
     /*
@@ -462,15 +491,10 @@ PyArray_ConcatenateArrays(int narrays, PyArrayObject **arrays, int axis)
  */
 NPY_NO_EXPORT PyArrayObject *
 PyArray_ConcatenateFlattenedArrays(int narrays, PyArrayObject **arrays,
-                                    NPY_ORDER order)
+                                    NPY_ORDER order, PyArrayObject *ret)
 {
-    PyTypeObject *subtype = &PyArray_Type;
-    double priority = NPY_PRIORITY;
     int iarrays;
-    npy_intp stride;
     npy_intp shape = 0;
-    PyArray_Descr *dtype = NULL;
-    PyArrayObject *ret = NULL;
     PyArrayObject_fields *sliding_view = NULL;
 
     if (narrays <= 0) {
@@ -494,36 +518,45 @@ PyArray_ConcatenateFlattenedArrays(int narrays, PyArrayObject **arrays,
         }
     }
 
-    /* Get the priority subtype for the array */
-    for (iarrays = 0; iarrays < narrays; ++iarrays) {
-        if (Py_TYPE(arrays[iarrays]) != subtype) {
-            double pr = PyArray_GetPriority((PyObject *)(arrays[iarrays]), 0.0);
-            if (pr > priority) {
-                priority = pr;
-                subtype = Py_TYPE(arrays[iarrays]);
-            }
+    if (ret != NULL) {
+        if (PyArray_NDIM(ret) != 1) {
+            PyErr_SetString(PyExc_ValueError,
+                            "Output array must be 1D");
+            return NULL;
+        }
+        if (shape != PyArray_SIZE(ret)) {
+            PyErr_SetString(PyExc_ValueError,
+                            "Output array is the wrong size");
+            return NULL;
         }
+        Py_INCREF(ret);
     }
+    else {
+        npy_intp stride;
 
-    /* Get the resulting dtype from combining all the arrays */
-    dtype = PyArray_ResultType(narrays, arrays, 0, NULL);
-    if (dtype == NULL) {
-        return NULL;
-    }
+        /* Get the priority subtype for the array */
+        PyTypeObject *subtype = PyArray_GetSubType(narrays, arrays);
 
-    stride = dtype->elsize;
+        /* Get the resulting dtype from combining all the arrays */
+        PyArray_Descr *dtype = PyArray_ResultType(narrays, arrays, 0, NULL);
+        if (dtype == NULL) {
+            return NULL;
+        }
 
-    /* Allocate the array for the result. This steals the 'dtype' reference. */
-    ret = (PyArrayObject *)PyArray_NewFromDescr(subtype,
-                                                    dtype,
-                                                    1,
-                                                    &shape,
-                                                    &stride,
-                                                    NULL,
-                                                    0,
-                                                    NULL);
-    if (ret == NULL) {
-        return NULL;
+        stride = dtype->elsize;
+
+        /* Allocate the array for the result. This steals the 'dtype' reference. */
+        ret = (PyArrayObject *)PyArray_NewFromDescr(subtype,
+                                                        dtype,
+                                                        1,
+                                                        &shape,
+                                                        &stride,
+                                                        NULL,
+                                                        0,
+                                                        NULL);
+        if (ret == NULL) {
+            return NULL;
+        }
     }
 
     /*
@@ -558,22 +591,11 @@ PyArray_ConcatenateFlattenedArrays(int narrays, PyArrayObject **arrays,
     return ret;
 }
 
-
-/*NUMPY_API
- * Concatenate
- *
- * Concatenate an arbitrary Python sequence into an array.
- * op is a python object supporting the sequence interface.
- * Its elements will be concatenated together to form a single
- * multidimensional array. If axis is NPY_MAXDIMS or bigger, then
- * each sequence object will be flattened before concatenation
-*/
 NPY_NO_EXPORT PyObject *
-PyArray_Concatenate(PyObject *op, int axis)
+PyArray_ConcatenateInto(PyObject *op, int axis, PyArrayObject *ret)
 {
     int iarrays, narrays;
     PyArrayObject **arrays;
-    PyArrayObject *ret;
 
     if (!PySequence_Check(op)) {
         PyErr_SetString(PyExc_TypeError,
@@ -606,10 +628,10 @@ PyArray_Concatenate(PyObject *op, int axis)
     }
 
     if (axis >= NPY_MAXDIMS) {
-        ret = PyArray_ConcatenateFlattenedArrays(narrays, arrays, NPY_CORDER);
+        ret = PyArray_ConcatenateFlattenedArrays(narrays, arrays, NPY_CORDER, ret);
     }
     else {
-        ret = PyArray_ConcatenateArrays(narrays, arrays, axis);
+        ret = PyArray_ConcatenateArrays(narrays, arrays, axis, ret);
     }
 
     for (iarrays = 0; iarrays < narrays; ++iarrays) {
@@ -629,6 +651,21 @@ PyArray_Concatenate(PyObject *op, int axis)
     return NULL;
 }
 
+/*NUMPY_API
+ * Concatenate
+ *
+ * Concatenate an arbitrary Python sequence into an array.
+ * op is a python object supporting the sequence interface.
+ * Its elements will be concatenated together to form a single
+ * multidimensional array. If axis is NPY_MAXDIMS or bigger, then
+ * each sequence object will be flattened before concatenation
+*/
+NPY_NO_EXPORT PyObject *
+PyArray_Concatenate(PyObject *op, int axis)
+{
+    return PyArray_ConcatenateInto(op, axis, NULL);
+}
+
 static int
 _signbit_set(PyArrayObject *arr)
 {
@@ -2159,14 +2196,24 @@ static PyObject *
 array_concatenate(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *kwds)
 {
     PyObject *a0;
+    PyObject *out = NULL;
     int axis = 0;
-    static char *kwlist[] = {"seq", "axis", NULL};
+    static char *kwlist[] = {"seq", "axis", "out", NULL};
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O&:concatenate", kwlist,
-                &a0, PyArray_AxisConverter, &axis)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O&O:concatenate", kwlist,
+                &a0, PyArray_AxisConverter, &axis, &out)) {
         return NULL;
     }
-    return PyArray_Concatenate(a0, axis);
+    if (out != NULL) {
+        if (out == Py_None) {
+            out = NULL;
+        }
+        else if (!PyArray_Check(out)) {
+            PyErr_SetString(PyExc_TypeError, "'out' must be an array");
+            return NULL;
+        }
+    }
+    return PyArray_ConcatenateInto(a0, axis, (PyArrayObject *)out);
 }
 
 static PyObject *