Nota
Las matrices admiten el protocolo iterador y se pueden iterar como las listas de Python. Ver el Indexación, corte e iteración en la guía de inicio rápido para ver ejemplos de uso básico. El resto de este documento presenta la nditer objeto y cubre un uso más avanzado.
El objeto iterador nditer, introducido en NumPy 1.6, proporciona muchas formas flexibles de visitar todos los elementos de una o más matrices de forma sistemática. Esta página presenta algunas formas básicas de usar el objeto para cálculos en matrices en Python, luego concluye con cómo se puede acelerar el ciclo interno en Cython. Desde la exposición de Python de nditer es un mapeo relativamente sencillo de la API del iterador de matrices de C, estas ideas también ayudarán a trabajar con la iteración de matrices desde C o C ++.
Iteración de matriz única
La tarea más básica que se puede realizar con el nditer es visitar todos los elementos de una matriz. Cada elemento se proporciona uno por uno utilizando la interfaz de iterador estándar de Python.
Ejemplo
>>> a = np.arange(6).reshape(2,3) >>> for x in np.nditer(a): ... print(x, end=' ') ... 0 1 2 3 4 5
Una cosa importante a tener en cuenta para esta iteración es que el orden se elige para que coincida con el diseño de memoria de la matriz en lugar de utilizar un orden estándar C o Fortran. Esto se hace para la eficiencia del acceso, reflejando la idea de que, por defecto, uno simplemente quiere visitar cada elemento sin preocuparse por un orden en particular. Podemos ver esto iterando sobre la transposición de nuestra matriz anterior, en comparación con tomar una copia de esa transposición en orden C.
Ejemplo
>>> a = np.arange(6).reshape(2,3) >>> for x in np.nditer(a.T): ... print(x, end=' ') ... 0 1 2 3 4 5
>>> for x in np.nditer(a.T.copy(order='C')): ... print(x, end=' ') ... 0 3 1 4 2 5
Los elementos de ambos a y a.T se recorren en el mismo orden, es decir, el orden en que se almacenan en la memoria, mientras que los elementos de a.T.copy(order=’C’) recibir visitas en un orden diferente porque se han colocado en un diseño de memoria diferente.
Control de orden de iteración
Hay ocasiones en las que es importante visitar los elementos de una matriz en un orden específico, independientemente del diseño de los elementos en la memoria. los nditer El objeto proporciona un order parámetro para controlar este aspecto de la iteración. El valor predeterminado, con el comportamiento descrito anteriormente, es order = ‘K’ para mantener el orden existente. Esto se puede anular con order = ‘C’ para el pedido C y order = ‘F’ para el pedido Fortran.
Ejemplo
>>> a = np.arange(6).reshape(2,3) >>> for x in np.nditer(a, order='F'): ... print(x, end=' ') ... 0 3 1 4 2 5 >>> for x in np.nditer(a.T, order='C'): ... print(x, end=' ') ... 0 3 1 4 2 5
Modificar valores de matriz
Por defecto, el nditer trata el operando de entrada como un objeto de solo lectura. Para poder modificar los elementos de la matriz, debe especificar el modo de lectura-escritura o de solo escritura usando el ‘readwrite’ o ‘writeonly’ banderas por operando.
El nditer generará matrices de búfer grabables que puede modificar. Sin embargo, debido a que el nditer debe copiar estos datos de búfer de nuevo a la matriz original una vez que finaliza la iteración, debe señalar cuándo finaliza la iteración, mediante uno de dos métodos. Puede:
- usó el nditer como administrador de contexto usando el
withdeclaración, y los datos temporales se volverán a escribir cuando se salga del contexto. - llamar al iterador
closeuna vez que haya terminado de iterar, lo que activará la escritura.
El nditer ya no se puede iterar una vez tampoco close se llama o se sale de su contexto.
Ejemplo
>>> a = np.arange(6).reshape(2,3) >>> a array([[0, 1, 2], [3, 4, 5]]) >>> with np.nditer(a, op_flags=['readwrite']) as it: ... for x in it: ... x[...] = 2 * x ... >>> a array([[ 0, 2, 4], [ 6, 8, 10]])
Si está escribiendo código que debe admitir versiones anteriores de numpy, tenga en cuenta que antes de la 1.15, nditer no era un administrador de contexto y no tenía un close método. En su lugar, se basó en el destructor para iniciar la escritura diferida del búfer.
Usando un lazo externo
En todos los ejemplos hasta ahora, los elementos de a son proporcionados por el iterador uno a la vez, porque toda la lógica de bucle es interna al iterador. Si bien esto es simple y conveniente, no es muy eficiente. Un mejor enfoque es mover el bucle más interno unidimensional a su código, externo al iterador. De esta manera, las operaciones vectorizadas de NumPy se pueden usar en partes más grandes de los elementos que se visitan.
los nditer intentará proporcionar fragmentos lo más grandes posible al bucle interior. Al forzar el orden ‘C’ y ‘F’, obtenemos diferentes tamaños de bucle externo. Este modo se habilita especificando un indicador de iterador.
Observe que con el valor predeterminado de mantener el orden de la memoria nativa, el iterador puede proporcionar un solo fragmento unidimensional, mientras que al forzar el orden de Fortran, debe proporcionar tres fragmentos de dos elementos cada uno.
Ejemplo
>>> a = np.arange(6).reshape(2,3) >>> for x in np.nditer(a, flags=['external_loop']): ... print(x, end=' ') ... [0 1 2 3 4 5]
>>> for x in np.nditer(a, flags=['external_loop'], order='F'): ... print(x, end=' ') ... [0 3] [1 4] [2 5]
Seguimiento de un índice o índice múltiple
Durante la iteración, es posible que desee utilizar el índice del elemento actual en un cálculo. Por ejemplo, es posible que desee visitar los elementos de una matriz en orden de memoria, pero utilice un orden C, un orden Fortran o un índice multidimensional para buscar valores en una matriz diferente.
El objeto iterador realiza un seguimiento del índice y se puede acceder a él a través del index o multi_index propiedades, dependiendo de lo solicitado. Los ejemplos siguientes muestran impresiones que demuestran la progresión del índice:
Ejemplo
>>> a = np.arange(6).reshape(2,3) >>> it = np.nditer(a, flags=['f_index']) >>> for x in it: ... print("%d <%d>" % (x, it.index), end=' ') ... 0 <0> 1 <2> 2 <4> 3 <1> 4 <3> 5 <5>
>>> it = np.nditer(a, flags=['multi_index']) >>> for x in it: ... print("%d <%s>" % (x, it.multi_index), end=' ') ... 0 <(0, 0)> 1 <(0, 1)> 2 <(0, 2)> 3 <(1, 0)> 4 <(1, 1)> 5 <(1, 2)>
>>> with np.nditer(a, flags=['multi_index'], op_flags=['writeonly']) as it: ... for x in it: ... x[...] = it.multi_index[1] - it.multi_index[0] ... >>> a array([[ 0, 1, 2], [-1, 0, 1]])
El seguimiento de un índice o de un índice múltiple es incompatible con el uso de un bucle externo, porque requiere un valor de índice diferente por elemento. Si intenta combinar estas banderas, la nditer objeto generará una excepción.
Ejemplo
>>> a = np.zeros((2,3)) >>> it = np.nditer(a, flags=['c_index', 'external_loop']) Traceback (most recent call last): File "<stdin>", line 1, in <module> ValueError: Iterator flag EXTERNAL_LOOP cannot be used if an index or multi-index is being tracked
Bucle alternativo y acceso a elementos
Para que sus propiedades sean más accesibles durante la iteración, nditer tiene una sintaxis alternativa para iterar, que funciona explícitamente con el objeto iterador en sí. Con esta construcción de bucle, se puede acceder al valor actual indexando en el iterador. Otras propiedades, como los índices de seguimiento, permanecen como antes. Los ejemplos siguientes producen resultados idénticos a los de la sección anterior.
Ejemplo
>>> a = np.arange(6).reshape(2,3) >>> it = np.nditer(a, flags=['f_index']) >>> while not it.finished: ... print("%d <%d>" % (it[0], it.index), end=' ') ... is_not_finished = it.iternext() ... 0 <0> 1 <2> 2 <4> 3 <1> 4 <3> 5 <5>
>>> it = np.nditer(a, flags=['multi_index']) >>> while not it.finished: ... print("%d <%s>" % (it[0], it.multi_index), end=' ') ... is_not_finished = it.iternext() ... 0 <(0, 0)> 1 <(0, 1)> 2 <(0, 2)> 3 <(1, 0)> 4 <(1, 1)> 5 <(1, 2)>
>>> with np.nditer(a, flags=['multi_index'], op_flags=['writeonly']) as it: ... while not it.finished: ... it[0] = it.multi_index[1] - it.multi_index[0] ... is_not_finished = it.iternext() ... >>> a array([[ 0, 1, 2], [-1, 0, 1]])
Almacenamiento en búfer de los elementos de la matriz
Al forzar un orden de iteración, observamos que la opción de bucle externo puede proporcionar los elementos en trozos más pequeños porque los elementos no se pueden visitar en el orden apropiado con un paso constante. Al escribir código C, esto generalmente está bien, sin embargo, en código Python puro, esto puede causar una reducción significativa en el rendimiento.
Al habilitar el modo de almacenamiento en búfer, los fragmentos proporcionados por el iterador al bucle interno se pueden agrandar, lo que reduce significativamente la sobrecarga del intérprete de Python. En el ejemplo que fuerza el orden de iteración de Fortran, el bucle interno puede ver todos los elementos de una vez cuando el almacenamiento en búfer está habilitado.
Ejemplo
>>> a = np.arange(6).reshape(2,3) >>> for x in np.nditer(a, flags=['external_loop'], order='F'): ... print(x, end=' ') ... [0 3] [1 4] [2 5]
>>> for x in np.nditer(a, flags=['external_loop','buffered'], order='F'): ... print(x, end=' ') ... [0 3 1 4 2 5]
Iterando como un tipo de datos específico
Hay ocasiones en las que es necesario tratar una matriz como un tipo de datos diferente al que está almacenado. Por ejemplo, es posible que desee hacer todos los cálculos en flotantes de 64 bits, incluso si las matrices que se manipulan son flotantes de 32 bits. Excepto cuando se escribe código C de bajo nivel, generalmente es mejor dejar que el iterador maneje la copia o el almacenamiento en búfer en lugar de convertir el tipo de datos usted mismo en el bucle interno.
Hay dos mecanismos que permiten hacer esto, las copias temporales y el modo de almacenamiento en búfer. Con copias temporales, se hace una copia de toda la matriz con el nuevo tipo de datos, luego se realiza una iteración en la copia. Se permite el acceso de escritura a través de un modo que actualiza la matriz original después de que se completa toda la iteración. El principal inconveniente de las copias temporales es que la copia temporal puede consumir una gran cantidad de memoria, especialmente si el tipo de datos de iteración tiene un tamaño de elemento mayor que el original.
El modo de almacenamiento en búfer mitiga el problema del uso de la memoria y es más amigable con el caché que hacer copias temporales. Excepto en casos especiales, donde se necesita toda la matriz a la vez fuera del iterador, se recomienda el almacenamiento en búfer en lugar de la copia temporal. Dentro de NumPy, ufuncs y otras funciones utilizan el almacenamiento en búfer para admitir entradas flexibles con una sobrecarga de memoria mínima.
En nuestros ejemplos, trataremos la matriz de entrada con un tipo de datos complejo, de modo que podamos sacar raíces cuadradas de números negativos. Sin habilitar las copias o el modo de almacenamiento en búfer, el iterador generará una excepción si el tipo de datos no coincide con precisión.
Ejemplo
>>> a = np.arange(6).reshape(2,3) - 3 >>> for x in np.nditer(a, op_dtypes=['complex128']): ... print(np.sqrt(x), end=' ') ... Traceback (most recent call last): File "<stdin>", line 1, in <module> TypeError: Iterator operand required copying or buffering, but neither copying nor buffering was enabled
En el modo de copia, ‘copiar’ se especifica como un indicador por operando. Esto se hace para proporcionar control por operando. El modo de almacenamiento en búfer se especifica como una bandera de iterador.
Ejemplo
>>> a = np.arange(6).reshape(2,3) - 3 >>> for x in np.nditer(a, op_flags=['readonly','copy'], ... op_dtypes=['complex128']): ... print(np.sqrt(x), end=' ') ... 1.7320508075688772j 1.4142135623730951j 1j 0j (1+0j) (1.4142135623730951+0j)
>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['complex128']): ... print(np.sqrt(x), end=' ') ... 1.7320508075688772j 1.4142135623730951j 1j 0j (1+0j) (1.4142135623730951+0j)
El iterador usa las reglas de conversión de NumPy para determinar si se permite una conversión específica. De forma predeterminada, impone la conversión “segura”. Esto significa, por ejemplo, que generará una excepción si intenta tratar una matriz flotante de 64 bits como una matriz flotante de 32 bits. En muchos casos, la regla ‘same_kind’ es la regla más razonable de usar, ya que permitirá la conversión de float de 64 a 32 bits, pero no de float a int o de complex a float.
Ejemplo
>>> a = np.arange(6.) >>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['float32']): ... print(x, end=' ') ... Traceback (most recent call last): File "<stdin>", line 1, in <module> TypeError: Iterator operand 0 dtype could not be cast from dtype('float64') to dtype('float32') according to the rule 'safe'
>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['float32'], ... casting='same_kind'): ... print(x, end=' ') ... 0.0 1.0 2.0 3.0 4.0 5.0
>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['int32'], casting='same_kind'): ... print(x, end=' ') ... Traceback (most recent call last): File "<stdin>", line 1, in <module> TypeError: Iterator operand 0 dtype could not be cast from dtype('float64') to dtype('int32') according to the rule 'same_kind'
Una cosa a tener en cuenta son las conversiones al tipo de datos original cuando se usa un operando de lectura-escritura o de solo escritura. Un caso común es implementar el bucle interno en términos de flotantes de 64 bits y usar la conversión ‘same_kind’ para permitir que también se procesen los otros tipos de punto flotante. Mientras que en el modo de solo lectura, se podría proporcionar una matriz de enteros, el modo de lectura y escritura generará una excepción porque la conversión de nuevo a la matriz violaría la regla de conversión.
Ejemplo
>>> a = np.arange(6) >>> for x in np.nditer(a, flags=['buffered'], op_flags=['readwrite'], ... op_dtypes=['float64'], casting='same_kind'): ... x[...] = x / 2.0 ... Traceback (most recent call last): File "<stdin>", line 2, in <module> TypeError: Iterator requested dtype could not be cast from dtype('float64') to dtype('int64'), the operand 0 dtype, according to the rule 'same_kind'
Iteración de matriz de difusión
NumPy tiene un conjunto de reglas para lidiar con matrices que tienen diferentes formas que se aplican siempre que las funciones toman múltiples operandos que combinan elementos. Se llama radiodifusión. los nditer object puede aplicar estas reglas por usted cuando necesite escribir dicha función.
Como ejemplo, imprimimos el resultado de transmitir una matriz de una y una de dos dimensiones juntas.
Ejemplo
>>> a = np.arange(3) >>> b = np.arange(6).reshape(2,3) >>> for x, y in np.nditer([a,b]): ... print("%d:%d" % (x,y), end=' ') ... 0:0 1:1 2:2 0:3 1:4 2:5
Cuando se produce un error de transmisión, el iterador genera una excepción que incluye las formas de entrada para ayudar a diagnosticar el problema.
Ejemplo
>>> a = np.arange(2) >>> b = np.arange(6).reshape(2,3) >>> for x, y in np.nditer([a,b]): ... print("%d:%d" % (x,y), end=' ') ... Traceback (most recent call last): ... ValueError: operands could not be broadcast together with shapes (2,) (2,3)
Matrices de salida asignadas por iteradores
Un caso común en las funciones de NumPy es tener salidas asignadas en función de la transmisión de la entrada y, además, tener un parámetro opcional llamado ‘out’ donde se colocará el resultado cuando se proporcione. los nditer El objeto proporciona un lenguaje conveniente que hace que sea muy fácil admitir este mecanismo.
Mostraremos cómo funciona esto creando una función. square que cuadra su entrada. Comencemos con una definición de función mínima que excluya el soporte de parámetros ‘out’.
Ejemplo
>>> def square(a): ... with np.nditer([a, None]) as it: ... for x, y in it: ... y[...] = x*x ... return it.operands[1] ... >>> square([1,2,3]) array([1, 4, 9])
Por defecto, el nditer utiliza los indicadores ‘asignar’ y ‘solo escritura’ para los operandos que se pasan como Ninguno. Esto significa que pudimos proporcionar solo los dos operandos al iterador y él manejó el resto.
Al agregar el parámetro ‘out’, tenemos que proporcionar explícitamente esos indicadores, porque si alguien pasa una matriz como ‘out’, el iterador se establecerá por defecto en ‘readonly’ y nuestro bucle interno fallará. La razón por la que ‘readonly’ es el valor predeterminado para las matrices de entrada es para evitar confusiones sobre la activación involuntaria de una operación de reducción. Si el valor predeterminado fuera ‘leer y escribir’, cualquier operación de transmisión también desencadenaría una reducción, un tema que se trata más adelante en este documento.
Ya que estamos en eso, introduzcamos también la bandera ‘no_broadcast’, que evitará que se difunda la salida. Esto es importante porque solo queremos un valor de entrada para cada salida. Agregar más de un valor de entrada es una operación de reducción que requiere un manejo especial. Ya generaría un error porque las reducciones deben habilitarse explícitamente en un indicador de iterador, pero el mensaje de error que resulta de deshabilitar la radiodifusión es mucho más comprensible para los usuarios finales. Para ver cómo generalizar la función cuadrada a una reducción, observe la función suma de cuadrados en la sección sobre Cython.
Para completar, también agregaremos los indicadores ‘external_loop’ y ‘buffer’, ya que estos son los que normalmente querrá por razones de rendimiento.
Ejemplo
>>> def square(a, out=None): ... it = np.nditer([a, out], ... flags = ['external_loop', 'buffered'], ... op_flags = [['readonly'], ... ['writeonly', 'allocate', 'no_broadcast']]) ... with it: ... for x, y in it: ... y[...] = x*x ... return it.operands[1] ...
>>> square([1,2,3]) array([1, 4, 9])
>>> b = np.zeros((3,)) >>> square([1,2,3], out=b) array([ 1., 4., 9.]) >>> b array([ 1., 4., 9.])
>>> square(np.arange(6).reshape(2,3), out=b) Traceback (most recent call last): ... ValueError: non-broadcastable output operand with shape (3,) doesn't match the broadcast shape (2,3)
Iteración externa del producto
Cualquier operación binaria puede extenderse a una operación de matriz en una forma de producto externo como en outer, y el nditer El objeto proporciona una manera de lograr esto mapeando explícitamente los ejes de los operandos. También es posible hacer esto con newaxis indexación, pero le mostraremos cómo usar directamente el nditer op_axes parámetro para lograr esto sin vistas intermedias.
Haremos un producto externo simple, colocando las dimensiones del primer operando antes de las dimensiones del segundo operando. los op_axes El parámetro necesita una lista de ejes para cada operando y proporciona una asignación de los ejes del iterador a los ejes del operando.
Suponga que el primer operando es unidimensional y el segundo operando es bidimensional. El iterador tendrá tres dimensiones, por lo que op_axes tendrá dos listas de 3 elementos. La primera lista selecciona el eje del primer operando, y es -1 para el resto de los ejes del iterador, con un resultado final de [0, -1, -1]. La segunda lista selecciona los dos ejes del segundo operando, pero no debe superponerse con los ejes seleccionados en el primer operando. Su lista es [-1, 0, 1]. El operando de salida se asigna a los ejes del iterador de la manera estándar, por lo que podemos proporcionar Ninguno en lugar de construir otra lista.
La operación en el bucle interno es una multiplicación sencilla. Todo lo que tiene que ver con el producto externo lo maneja la configuración del iterador.
Ejemplo
>>> a = np.arange(3) >>> b = np.arange(8).reshape(2,4) >>> it = np.nditer([a, b, None], flags=['external_loop'], ... op_axes=[[0, -1, -1], [-1, 0, 1], None]) >>> with it: ... for x, y, z in it: ... z[...] = x*y ... result = it.operands[2] # same as z ... >>> result array([[[ 0, 0, 0, 0], [ 0, 0, 0, 0]], [[ 0, 1, 2, 3], [ 4, 5, 6, 7]], [[ 0, 2, 4, 6], [ 8, 10, 12, 14]]])
Tenga en cuenta que una vez cerrado el iterador no podemos acceder operands y debe utilizar una referencia creada dentro del administrador de contexto.
Iteración de reducción
Siempre que un operando grabable tiene menos elementos que el espacio de iteración completo, ese operando sufre una reducción. los nditer El objeto requiere que cualquier operando de reducción sea marcado como lectura-escritura, y solo permite reducciones cuando ‘reduce_ok’ se proporciona como un indicador de iterador.
Para obtener un ejemplo simple, considere tomar la suma de todos los elementos de una matriz.
Ejemplo
>>> a = np.arange(24).reshape(2,3,4) >>> b = np.array(0) >>> with np.nditer([a, b], flags=['reduce_ok'], ... op_flags=[['readonly'], ['readwrite']]) as it: ... for x,y in it: ... y[...] += x ... >>> b array(276) >>> np.sum(a) 276
Las cosas son un poco más complicadas cuando se combinan la reducción y los operandos asignados. Antes de que se inicie la iteración, cualquier operando de reducción debe inicializarse a sus valores iniciales. Así es como podemos hacer esto, tomando sumas a lo largo del último eje de a.
Ejemplo
>>> a = np.arange(24).reshape(2,3,4) >>> it = np.nditer([a, None], flags=['reduce_ok'], ... op_flags=[['readonly'], ['readwrite', 'allocate']], ... op_axes=[None, [0,1,-1]]) >>> with it: ... it.operands[1][...] = 0 ... for x, y in it: ... y[...] += x ... result = it.operands[1] ... >>> result array([[ 6, 22, 38], [54, 70, 86]]) >>> np.sum(a, axis=2) array([[ 6, 22, 38], [54, 70, 86]])
Para realizar la reducción con búfer se requiere otro ajuste durante la configuración. Normalmente, la construcción del iterador implica copiar el primer búfer de datos de las matrices legibles al búfer. Cualquier operando de reducción es legible, por lo que puede leerse en un búfer. Desafortunadamente, la inicialización del operando después de que se complete esta operación de almacenamiento en búfer no se reflejará en el búfer con el que comienza la iteración y se producirán resultados basura.
El indicador de iterador “delay_bufalloc” está ahí para permitir que los operandos de reducción asignados al iterador existan junto con el almacenamiento en búfer. Cuando se establece esta bandera, el iterador dejará sus búferes sin inicializar hasta que reciba un reinicio, después de lo cual estará listo para la iteración regular. Así es como se ve el ejemplo anterior si también habilitamos el almacenamiento en búfer.
Ejemplo
>>> a = np.arange(24).reshape(2,3,4) >>> it = np.nditer([a, None], flags=['reduce_ok', ... 'buffered', 'delay_bufalloc'], ... op_flags=[['readonly'], ['readwrite', 'allocate']], ... op_axes=[None, [0,1,-1]]) >>> with it: ... it.operands[1][...] = 0 ... it.reset() ... for x, y in it: ... y[...] += x ... result = it.operands[1] ... >>> result array([[ 6, 22, 38], [54, 70, 86]])
Poner el bucle interno en Cython
Aquellos que desean un rendimiento realmente bueno de sus operaciones de bajo nivel deben considerar seriamente el uso directo de la API de iteración proporcionada en C, pero para aquellos que no se sienten cómodos con C o C ++, Cython es un buen término medio con compensaciones de rendimiento razonables. Para el nditer object, esto significa dejar que el iterador se encargue de la transmisión, la conversión de dtype y el almacenamiento en búfer, mientras le da el bucle interno a Cython.
Para nuestro ejemplo, crearemos una función de suma de cuadrados. Para empezar, implementemos esta función en Python sencillo. Queremos admitir un parámetro de ‘eje’ similar al numpy sum función, por lo que necesitaremos construir una lista para la op_axes parámetro. Así es como se ve esto.
Ejemplo
>>> def axis_to_axeslist(axis, ndim): ... if axis is None: ... return [-1] * ndim ... else: ... if type(axis) is not tuple: ... axis = (axis,) ... axeslist = [1] * ndim ... for i in axis: ... axeslist[i] = -1 ... ax = 0 ... for i in range(ndim): ... if axeslist[i] != -1: ... axeslist[i] = ax ... ax += 1 ... return axeslist ... >>> def sum_squares_py(arr, axis=None, out=None): ... axeslist = axis_to_axeslist(axis, arr.ndim) ... it = np.nditer([arr, out], flags=['reduce_ok', ... 'buffered', 'delay_bufalloc'], ... op_flags=[['readonly'], ['readwrite', 'allocate']], ... op_axes=[None, axeslist], ... op_dtypes=['float64', 'float64']) ... with it: ... it.operands[1][...] = 0 ... it.reset() ... for x, y in it: ... y[...] += x*x ... return it.operands[1] ... >>> a = np.arange(6).reshape(2,3) >>> sum_squares_py(a) array(55.0) >>> sum_squares_py(a, axis=-1) array([ 5., 50.])
Para Cython-ize esta función, reemplazamos el ciclo interno (y[…] + = x * x) con código Cython especializado para el tipo d float64. Con el indicador ‘external_loop’ habilitado, las matrices proporcionadas al bucle interno siempre serán unidimensionales, por lo que es necesario realizar muy pocas comprobaciones.
Aquí está la lista de sum_squares.pyx:
import numpy as np cimport numpy as np cimport cython def axis_to_axeslist(axis, ndim): if axis is None: return [-1] * ndim else: if type(axis) is not tuple: axis = (axis,) axeslist = [1] * ndim for i in axis: axeslist[i] = -1 ax = 0 for i in range(ndim): if axeslist[i] != -1: axeslist[i] = ax ax += 1 return axeslist @cython.boundscheck(False) def sum_squares_cy(arr, axis=None, out=None): cdef np.ndarray[double] x cdef np.ndarray[double] y cdef int size cdef double value axeslist = axis_to_axeslist(axis, arr.ndim) it = np.nditer([arr, out], flags=['reduce_ok', 'external_loop', 'buffered', 'delay_bufalloc'], op_flags=[['readonly'], ['readwrite', 'allocate']], op_axes=[None, axeslist], op_dtypes=['float64', 'float64']) with it: it.operands[1][...] = 0 it.reset() for xarr, yarr in it: x = xarr y = yarr size = x.shape[0] for i in range(size): value = x[i] y[i] = y[i] + value * value return it.operands[1]
En esta máquina, compilar el archivo .pyx en un módulo se veía como el siguiente, pero es posible que tenga que encontrar algunos tutoriales de Cython que le indiquen los detalles de la configuración de su sistema:
$ cython sum_squares.pyx $ gcc -shared -pthread -fPIC -fwrapv -O2 -Wall -I/usr/include/python2.7 -fno-strict-aliasing -o sum_squares.so sum_squares.c
Ejecutar esto desde el intérprete de Python produce las mismas respuestas que hizo nuestro código nativo de Python / NumPy.
Ejemplo
>>> from sum_squares import sum_squares_cy >>> a = np.arange(6).reshape(2,3) >>> sum_squares_cy(a) array(55.0) >>> sum_squares_cy(a, axis=-1) array([ 5., 50.])
Hacer un poco de tiempo en IPython muestra que la sobrecarga reducida y la asignación de memoria del bucle interno de Cython está proporcionando una aceleración muy agradable tanto sobre el código Python sencillo como sobre una expresión que usa la función de suma incorporada de NumPy:
>>> a = np.random.rand(1000,1000) >>> timeit sum_squares_py(a, axis=-1) 10 loops, best of 3: 37.1 ms per loop >>> timeit np.sum(a*a, axis=-1) 10 loops, best of 3: 20.9 ms per loop >>> timeit sum_squares_cy(a, axis=-1) 100 loops, best of 3: 11.8 ms per loop >>> np.all(sum_squares_cy(a, axis=-1) == np.sum(a*a, axis=-1)) True >>> np.all(sum_squares_py(a, axis=-1) == np.sum(a*a, axis=-1)) True