多维数组的就地改组

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]

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

非常感谢！

Answer 1

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

• 指针数组存储值的内存地址沿axis
• 您的算法与修改 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

Answer 2

以下算法基于切片，不进行任何复制，它应该适用于任何np.ndarray。主要步骤是：

• np.ndindex() 用于遍历不同的多维索引，不包括属于您要洗牌的轴的索引
• 应用您已经为一维情况开发的 shuffle。

代码：

def shuffleND(np.ndarray x, axis=-1):
    cdef np.ndarray[long long, ndim=1] idx
    cdef unsigned int i, j, n, m
    if axis==-1:
        axis = x.ndim-1
    all_shape = list(np.shape(x))
    shape = all_shape[:]
    shape.pop(axis)
    for slices in np.ndindex(*shape):
        slices = list(slices)
        axis_slice = slices[:]
        axis_slice.insert(axis, slice(None))
        idx = np.where(~np.isnan(x[tuple(axis_slice)]))[0]
        for i in range(idx.shape[0]-1, 0, -1):
            j = randint(i+1)
            n, m = idx[i], idx[j]
            slice1 = slices[:]
            slice1.insert(axis, n)
            slice2 = slices[:]
            slice2.insert(axis, m)
            slice1 = tuple(slice1)
            slice2 = tuple(slice2)
            x[slice1], x[slice2] = x[slice2], x[slice1]
    return x