python - 使用 C 扩展 numpy 时如何考虑列连续数组

Question

我有一个 C 函数来规范化对数空间中数组的行(这可以防止数值下溢)。

我的C函数原型(prototype)如下:

void normalize_logspace_matrix(size_t nrow, size_t ncol, double* mat);

您可以看到它接受一个指向数组的指针并就地修改它。 C 代码当然假定数据保存为 C 连续数组，即连续行。

我使用 Cython 将函数包装如下(省略了导入和 cdef extern from):

def normalize_logspace(np.ndarray[np.double_t, ndim=2] mat):
    cdef Py_ssize_t n, d
    n = mat.shape[0]
    d = mat.shape[1]
    normalize_logspace_matrix(n, d, <double*> mat.data)
    return mat

大多数时候 numpy 数组是行连续的，并且该函数工作正常。但是，如果先前已转置 numpy 数组，则不会复制数据，而只会返回数据的新 View 。在这种情况下，我的函数失败了，因为数组不再是行连续的。

我可以通过将数组定义为具有 Fortran 连续顺序来解决这个问题，这样在转置之后它将是 C 连续的:

A = np.array([some_func(d) for d in range(D)], order='F').T
A = normalize_logspace(A)

显然，这很容易出错，用户必须注意数组的顺序是否正确，而在 Python 中用户不需要关心这一点。

让我同时处理行和列连续数组的最佳方法是什么？我假设在 Cython 中进行某种数组顺序检查是可行的方法。当然，我更喜欢不需要将数据复制到新数组中的解决方案，但我几乎认为这是必要的。

Answer 1

如果您想在不进行复制的情况下支持 C 和 Fortran 顺序的数组，您的 C 函数需要足够灵活以支持这两种顺序。这可以通过将 NumPy 数组的步长传递给 C 函数来实现:将原型(prototype)更改为

void normalize_logspace_matrix(size_t nrow, size_t ncol, 
                               size_t rowstride, size_t colstride,
                               double* mat);

和 Cython 调用

def normalize_logspace(np.ndarray[np.double_t, ndim=2] mat):
    cdef Py_ssize_t n, d, rowstride, colstride
    n = mat.shape[0]
    d = mat.shape[1]
    rowstride = mat.strides[0] // mat.itemsize
    colstride = mat.strides[1] // mat.itemsize
    normalize_logspace_matrix(n, d, rowstride, colstride, <double*> mat.data)
    return mat

然后，用 mat[row*rowstride + col*colstride] 替换 C 代码中出现的每个 mat[row*ncol + col]。