ENH: core: Add PyArray_NewLikeArray function

This function implements the numpy.empty_like semantics, but supports
the new NPY_KEEPORDER enumeration value as well as switching to a
different data type.

doc/source/reference/c-api.array.rst

@@ -158,6 +158,21 @@ From scratch
     *dims* and *strides* are copied into newly allocated dimension and
     strides arrays for the new array object.
 
+.. cfunction:: PyObject* PyArray_NewLikeArray(PyArrayObject* prototype, NPY_ORDER order, PyArray_Descr* descr)
+
+   This function steals a reference to *descr* if it is not NULL.
+
+   This array creation routine allows for the convenient creation of
+   a new array matching an existing array's shapes and memory layout,
+   possibly changing the layout and/or data type.
+
+   When *order* is NPY_ANYORDER, the result order is NPY_FORTRANORDER if
+   *prototype* is a fortran array, NPY_CORDER otherwise.  When *order* is
+   NPY_KEEPORDER, the result order matches that of *prototype*, even
+   when the axes of *prototype* aren't in C or Fortran order.
+
+   If *descr* is NULL, the data type of *prototype* is used.
+
 .. cfunction:: PyObject* PyArray_New(PyTypeObject* subtype, int nd, npy_intp* dims, int type_num, npy_intp* strides, void* data, int itemsize, int flags, PyObject* obj)
 
     This is similar to :cfunc:`PyArray_DescrNew` (...) except you

numpy/core/code_generators/numpy_api.py

@@ -308,6 +308,7 @@
     'PyArray_CanCastTypeTo':                271,
     'PyArray_EinsteinSum':                  272,
     'PyArray_FillWithZero':                 273,
+    'PyArray_NewLikeArray':                 274,
 }
 
 ufunc_types_api = {

numpy/core/include/numpy/ndarraytypes.h

@@ -825,6 +825,8 @@ typedef int (PyArray_FinalizeFunc)(PyArrayObject *, PyObject *);
 #define PyArray_ISWRITEABLE(m) PyArray_CHKFLAGS(m, NPY_WRITEABLE)
 #define PyArray_ISALIGNED(m) PyArray_CHKFLAGS(m, NPY_ALIGNED)
 
+#define PyArray_IS_C_CONTIGUOUS(m) PyArray_CHKFLAGS(m, NPY_C_CONTIGUOUS)
+#define PyArray_IS_F_CONTIGUOUS(m) PyArray_CHKFLAGS(m, NPY_F_CONTIGUOUS)
 
 #if NPY_ALLOW_THREADS
 #define NPY_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS

numpy/core/src/multiarray/convert.c

@@ -429,23 +429,13 @@ PyArray_FillWithZero(PyArrayObject *a)
  * Copy an array.
  */
 NPY_NO_EXPORT PyObject *
-PyArray_NewCopy(PyArrayObject *m1, NPY_ORDER fortran)
+PyArray_NewCopy(PyArrayObject *m1, NPY_ORDER order)
 {
-    PyArrayObject *ret;
-    if (fortran == PyArray_ANYORDER)
-        fortran = PyArray_ISFORTRAN(m1);
-
-    Py_INCREF(m1->descr);
-    ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(m1),
-                                                m1->descr,
-                                                m1->nd,
-                                                m1->dimensions,
-                                                NULL, NULL,
-                                                fortran,
-                                                (PyObject *)m1);
+    PyArrayObject *ret = PyArray_NewLikeArray(m1, order, NULL);
     if (ret == NULL) {
         return NULL;
     }
+
     if (PyArray_CopyInto(ret, m1) == -1) {
         Py_DECREF(ret);
         return NULL;

numpy/core/src/multiarray/ctors.c

@@ -498,6 +498,8 @@ int _arrays_overlap(PyArrayObject *arr1, PyArrayObject *arr2)
  *
  * Instead of trying to be fancy, we simply check for overlap and make
  * a temporary copy when one exists.
+ *
+ * Returns 0 on success, negative on failure.
  */
 NPY_NO_EXPORT int
 PyArray_MoveInto(PyArrayObject *dst, PyArrayObject *src)
@@ -522,17 +524,8 @@ PyArray_MoveInto(PyArrayObject *dst, PyArrayObject *src)
 
         /*
          * Allocate a temporary copy array.
-         * TODO: For efficiency, this should have a memory ordering
-         *       matching 'dst', even if 'dst' has its axes arbitrarily
-         *       scrambled.  A function to allocate this array needs to
-         *       be created.
          */
-        Py_INCREF(PyArray_DESCR(dst));
-        tmp = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type,
-                        PyArray_DESCR(dst),
-                        PyArray_NDIM(dst), PyArray_DIMS(dst), NULL,
-                        NULL, PyArray_ISFORTRAN(dst) ? NPY_F_CONTIGUOUS : 0,
-                        NULL);
+        tmp = (PyArrayObject *)PyArray_NewLikeArray(dst, NPY_KEEPORDER, NULL);
         if (tmp == NULL) {
             return -1;
         }
@@ -1235,6 +1228,135 @@ PyArray_NewFromDescr(PyTypeObject *subtype, PyArray_Descr *descr, int nd,
     return NULL;
 }
 
+typedef struct {
+    npy_intp perm, stride;
+} _npy_stride_sort_item;
+
+/*
+ * Sorts items so stride is descending, because C-order
+ * is the default in the face of ambiguity.
+ */
+int _npy_stride_sort_item_comparator(const void *a, const void *b)
+{
+    npy_intp astride = ((_npy_stride_sort_item *)a)->stride,
+            bstride = ((_npy_stride_sort_item *)b)->stride;
+
+    if (astride > bstride) {
+        return -1;
+    }
+    else if (astride == bstride) {
+        /*
+         * Make the qsort stable by next comparing the perm order.
+         * (Note that two perm entries will never be equal)
+         */
+        npy_intp aperm = ((_npy_stride_sort_item *)a)->perm,
+                bperm = ((_npy_stride_sort_item *)b)->perm;
+        return (aperm < bperm) ? -1 : 1;
+    }
+    else {
+        return 1;
+    }
+}
+
+/*NUMPY_API
+ * Creates a new array with the same shape as the provided one,
+ * with possible memory layout order and data type changes.
+ *
+ * prototype - The array the new one should be like.
+ * order     - NPY_CORDER - C-contiguous result.
+ *             NPY_FORTRANORDER - Fortran-contiguous result.
+ *             NPY_ANYORDER - Fortran if prototype is Fortran, C otherwise.
+ *             NPY_KEEPORDER - Keeps the axis ordering of prototype.
+ * dtype     - If not NULL, overrides the data type of the result.
+ *
+ * NOTE: If dtype is not NULL, steals the dtype reference.
+ */
+NPY_NO_EXPORT PyObject *
+PyArray_NewLikeArray(PyArrayObject *prototype, NPY_ORDER order,
+                     PyArray_Descr *dtype)
+{
+    PyObject *ret = NULL;
+    int ndim = PyArray_NDIM(prototype);
+
+    /* If no override data type, use the one from the prototype */
+    if (dtype == NULL) {
+        dtype = PyArray_DESCR(prototype);
+        Py_INCREF(dtype);
+    }
+
+    /* Handle ANYORDER and simple KEEPORDER cases */
+    switch (order) {
+        case NPY_ANYORDER:
+            order = PyArray_ISFORTRAN(prototype) ?
+                                    NPY_FORTRANORDER : NPY_CORDER;
+            break;
+        case NPY_KEEPORDER:
+            if (PyArray_IS_C_CONTIGUOUS(prototype) || ndim <= 1) {
+                order = NPY_CORDER;
+                break;
+            }
+            else if (PyArray_IS_F_CONTIGUOUS(prototype)) {
+                order = NPY_FORTRANORDER;
+                break;
+            }
+            break;
+        default:
+            break;
+    }
+
+    /* If it's not KEEPORDER, this is simple */
+    if (order != NPY_KEEPORDER) {
+        ret = PyArray_NewFromDescr(Py_TYPE(prototype),
+                                            dtype,
+                                            ndim,
+                                            PyArray_DIMS(prototype),
+                                            NULL,
+                                            NULL,
+                                            order,
+                                            (PyObject *)prototype);
+    }
+    /* KEEPORDER needs some analysis of the strides */
+    else {
+        npy_intp strides[NPY_MAXDIMS], stride;
+        npy_intp *shape = PyArray_DIMS(prototype);
+        _npy_stride_sort_item sortstrides[NPY_MAXDIMS];
+        int i, ndim = PyArray_NDIM(prototype);
+
+        /* Set up the permutation and absolute value of strides */
+        for (i = 0; i < ndim; ++i) {
+            sortstrides[i].perm = i;
+            sortstrides[i].stride = PyArray_STRIDE(prototype, i);
+            if (sortstrides[i].stride < 0) {
+                sortstrides[i].stride = -sortstrides[i].stride;
+            }
+        }
+
+        /* Sort them */
+        qsort(sortstrides, ndim, sizeof(_npy_stride_sort_item),
+                                        &_npy_stride_sort_item_comparator);
+
+        /* Build the new strides */
+        stride = dtype->elsize;
+        for (i = ndim-1; i >= 0; --i) {
+            npy_intp i_perm = sortstrides[i].perm;
+            strides[i_perm] = stride;
+            stride *= shape[i_perm];
+        }
+
+        /* Finally, allocate the array */
+        ret = PyArray_NewFromDescr(Py_TYPE(prototype),
+                                            dtype,
+                                            ndim,
+                                            shape,
+                                            strides,
+                                            NULL,
+                                            0,
+                                            (PyObject *)prototype);
+    }
+
+    return ret;
+}
+
 /*NUMPY_API
  * Generic new array creation routine.
  */