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對象。