n维数组上的cython元组索引

Question

我正在尝试将一些 工作 现有的 numpy/python 代码移植到 cython

我遇到的一个问题是我不能在 cython 中对多维数组使用元组索引，而在 python/numpy 中它确实有效。

这是一个简单的mwe：

cython_indexing.pyx

# cython: boundscheck=False
# cython: wraparound=False
def loop(int axis, double[:, :, :] a, double[:, :, :] b):
    cdef:
        int k, j, i
        tuple q, qp1

    for k in range(a.shape[0]):
        for j in range(a.shape[1]):
            for i in range(a.shape[2]):
                q = (k, j, i)
                if axis == 0:
                    qp1 = (k + 1, j, i)
                elif axis == 1:
                    qp1 = (k, j + 1, i)
                elif axis == 2:
                    qp1 = (k, j, i + 1)

                # ...
                # some other operations
                # with heavy reuse of q, qp1
                # ...

                a[q] = a[q] - (b[qp1] - b[q])

test_indexing.py

import pyximport; pyximport.install()
import numpy as np
from cython_indexing import loop

a = np.arange(27).astype('float').reshape(3, 3, 3)
b = a**2

for axis in (0, 1, 2):
    loop(axis, a, b)

此示例在 b[qp1] - b[q] 的编译时引发错误：

Invalid operand types for '-' (double[:, :]; double[:, :])

是否有任何简单的解决方案不涉及更改代码架构？

Answer 1

根本问题是 Cython 在编译时不知道元组有多大，因此它无法在编译时明智地进行数组索引 - 它不知道它必须返回的数组有多少维有。 （看起来它只是被混淆了，但即使它确实有效，也必须采用“通用 Python __getitem__ 代码路径，所以你不会加快速度）。

您可以进行两个（不太困难的）更改。第一个是做我认为你说的时候要避免的事情

任何简单的解决方案不涉及更改代码架构

是使用 3 个整数而不是一个元组：

cdef:
   int q0, q1, q2, qp1_0, qp1_1, qp1_2

# ....
a[q0,q1,q2] = a[q0,q1,q2] - (b[qp1_0,qp1_1,qp1_2] - b[q0,q1,q2])

第二个是不使用 Cython 的“typed-memoryview”接口，让a 和b 无类型：

def loop(int axis, a, b):

这将使索引能够与元组一起工作（就像在纯 Python 中一样），但不会比纯 Python 快得多。

不幸的是，这是一种权衡：如果您想要更快的速度，那么您必须避免使用 Python 对象，例如元组。