python - 多维数组的就地改组

Question

我正在尝试在 Cython 中实现一个 NaN 安全的洗牌程序，它可以沿着任意维度的多维矩阵的多个轴洗牌。

在一维矩阵的简单情况下，可以使用 Fisher–Yates 算法简单地混洗所有具有非 NaN 值的索引:

def shuffle1D(np.ndarray[double, ndim=1] x):
    cdef np.ndarray[long, ndim=1] idx = np.where(~np.isnan(x))[0]
    cdef unsigned int i,j,n,m

    randint = np.random.randint
    for i in xrange(len(idx)-1, 0, -1):
        j = randint(i+1)
        n,m = idx[i], idx[j]
        x[n], x[m] = x[m], x[n]

我想扩展此算法以处理大型多维数组而无需重新整形(这会触发此处未考虑的更复杂情况的副本)。为此，我需要摆脱固定的输入维度，这对于 Cython 中的 numpy 数组和 memoryviews 似乎都是不可能的。有解决方法吗？

非常感谢!

Answer 1

感谢@Veedrac 的评论，此答案使用了更多 Cython 功能。

• 指针数组存储沿轴的值的内存地址
• 您的算法经过修改后使用 that checks for nan values ，防止它们被分类
• 它不会为 C 有序数组创建副本。对于 Fortran 有序数组，ravel() 命令将返回一个副本。这可以通过创建另一个双指针数组来携带 x 的值来改进，可能会有一些缓存惩罚......

此代码至少比另一个基于切片的代码快一个数量级。

from libc.stdlib cimport malloc, free

cimport numpy as np
import numpy as np
from numpy.random import randint

cdef extern from "numpy/npy_math.h":
    bint npy_isnan(double x)

def shuffleND(x, int axis=-1):
    cdef np.ndarray[double, ndim=1] v # view of x
    cdef np.ndarray[int, ndim=1] strides
    cdef int i, j
    cdef int num_axis, pos, stride
    cdef double tmp
    cdef double **v_axis

    if axis==-1:
        axis = x.ndim-1

    shape = list(x.shape)
    num_axis = shape.pop(axis)

    v_axis = <double **>malloc(num_axis*sizeof(double *))
    for i in range(num_axis):
        v_axis[i] = <double *>malloc(1*sizeof(double))

    try:
        tmp_strides = [s//x.itemsize for s in x.strides]
        stride = tmp_strides.pop(axis)
        strides = np.array(tmp_strides, dtype=np.int32)
        v = x.ravel()
        for indices in np.ndindex(*shape):
            pos = (strides*indices).sum()
            for i in range(num_axis):
                v_axis[i] = &v[pos + i*stride]
            for i in range(num_axis-1, 0, -1):
                j = randint(i+1)
                if npy_isnan(v_axis[i][0]) or npy_isnan(v_axis[j][0]):
                    continue
                tmp = v_axis[i][0]
                v_axis[i][0] = v_axis[j][0]
                v_axis[j][0] = tmp
    finally:
        free(v_axis)

    return x