Cython кортеж индексирование на п-мерный массив

Question

Я пытаюсь порта некоторые работает существующий NumPy/питон код Cython

Одна из проблем, я ударил в том, что я не могу использовать кортеж индексации для многомерного массива в Cython, тогда как в python/numpy он работает.

здесь просто м.в.э.:

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 ints вместо кортежа:

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]) 

Второй будет не использует интерфейс «типизированных memoryview» Cython, и позволить a и b быть нетипизированным:

def loop(int axis, a, b): 

Это позволило бы проиндексировать работу с кортежами (как в чистом Python), но не будет значительно быстрее, чем чистый Python.

К сожалению, это компромисс: если вам нужна более высокая скорость, вам нужно избегать использования объектов Python, таких как кортежи.