Accediendo a las columnas de la matriz de registro NumPy en Cython

Question

Soy un programador de Python relativamente experimentado, pero no he escrito ninguna C en mucho tiempo y estoy tratando de entender Cython. Intento escribir una función Cython que operará en una columna de un recarray NumPy.

El código que tengo hasta ahora está debajo.

recarray_func.pyx:

import numpy as np 
cimport numpy as np 

cdef packed struct rec_cell0: 
    np.float32_t f0 
    np.int64_t i0, i1, i2 

def sum(np.ndarray[rec_cell0, ndim=1] recarray): 
    cdef Py_ssize_t i 
    cdef rec_cell0 *cell 
    cdef np.float32_t running_sum = 0 

    for i in range(recarray.shape[0]): 
     cell = &recarray[i] 
     running_sum += cell.f0 
    return running_sum 

En el símbolo del intérprete:

array = np.recarray((100,), names=['f0', 'i0', 'i1', 'i2'], 
          formats=['f4', 'i8', 'i8', 'i8']) 
recarray_func.sum(array) 

Esto simplemente resume la columna f0 de la recArray. Se compila y se ejecuta sin problemas.

Mi pregunta es, ¿cómo podría modificar esto para que pueda operar en cualquier columna? En el ejemplo anterior, la columna a suma está codificada y se accede a ella a través de la notación de puntos. ¿Es posible cambiar la función para que la columna a la suma se pase como un parámetro?

Answer 1

Lo que usted necesita requiere tipeo débil, que C no tiene. Si todos sus tipos de registro son iguales, es posible que pueda realizar algo como: (descargo de responsabilidad, no tengo Cython en esta máquina, así que estoy codificando a ciegas).

import numpy as np 
cimport numpy as np 

cdef packed struct rec_cell0: 
    np.float32_t f0 
    np.int64_t i0, i1, i2 

def sum(np.ndarray[rec_cell0, ndim=1] recarray, colname): 
    cdef Py_ssize_t i 
    cdef rec_cell0 *cell 
    cdef np.float32_t running_sum = 0 

    loc = recarray.dtype.fields[colname][1] 

    for i in range(recarray.shape[0]): 
     cell = &recarray[i] 
     running_sum += *(int *)(&cell+loc); 
    return running_sum 

Answer 2

creo que esto debería ser posible utilizando Cython de memoryviews. Algo a lo largo de estas líneas debe trabajar (código no probado):

import numpy as np 
cimport numpy as np 

cdef packed struct rec_cell0: 
    np.float32_t f0 
    np.int64_t i0, i1, i2 

def sum(rec_cell0[:] recview): 
    cdef Py_ssize_t i 
    cdef np.float32_t running_sum = 0 

    for i in range(recview.shape[0]): 
     running_sum += recview[i].f0 
    return running_sum 

velocidad, probablemente, se puede aumentar al asegurar que la matriz de registro que pasa a Cython es contigua. En el lado de python (llamada), puede usar np.require, mientras que la firma de la función debe cambiar a rec_cell0[::1] recview para indicar que se puede suponer que la matriz es contigua. Y como siempre, una vez que el código ha sido probado, apagar el boundscheck, wraparound y nonecheckcompiler directives en Cython probablemente mejorará aún más la velocidad.