Cómo acelerar tu código Python con numba

Introducción

En este artículo vamos a hacer un repaso exhaustivo de cómo acelerar sustancialmente tu código Python usando numba. Ya hablamos sobre la primera versión de numba en el blog, allá por 2012, pero ha habido importantes cambios desde entonces y la herramienta ha cambiado muchísimo. Recientemente Continuum publicó numba 0.17 con una nueva documentación mucho más fácil de seguir, pero aun así no siempre queda claro cuál es el camino para hacer que funcione, como quedó patente con el artículo sobre Cython de Kiko. Por ello, en este artículo voy a aclarar qué puede y qué no puede hacer numba, cómo sacarle partido y voy a detallar un par de ejemplos exitosos que he producido en los últimos meses.

Hablando de las nuevas versiones de numba, en su web podéis ver una evolución temporal del rendimiento de algunas tareas que utiliza asv para la visualización.

Entendiendo numba: el modo nopython

Como podemos leer en la documentación, numba tiene dos modos de funcionamiento básicos: el modo object y el modo nopython.

  • El modo object genera código que gestiona todas las variables como objetos de Python y utiliza la API C de Python para operar con ellas. En general en este modo no habrá ganancias de rendimiento (e incluso puede ir más lento), con lo cual mi recomendación personal es directamente no utilizarlo. Hay casos en los que numba puede detectar los bucles y optimizarlos individualmente (loop-jitting), pero no le voy a prestar mucha atención a esto.
  • El modo nopython genera código independiente de la API C de Python. Esto tiene la desventaja de que no podemos usar todas las características del lenguaje, pero tiene un efecto significativo en el rendimiento. Otra de las restricciones es que no se puede reservar memoria para objetos nuevos.

Por defecto numba usa el modo nopython siempre que puede, y si no pasa a modo object. Nosotros vamos a forzar el modo nopython (o «modo estricto» como me gusta llamarlo) porque es la única forma de aprovechar el potencial de numba.

Ámbito de aplicación

El problema del modo nopython es que los mensajes de error son totalmente inservibles en la mayoría de los casos, así que antes de lanzarnos a compilar funciones con numba conviene hacer un repaso de qué no podemos hacer para anticipar la mejor forma de programar nuestro código. Podéis consultar en la documentación el subconjunto de Python soportado por numba en modo nopython, y ya os aviso que, al menos de momento, no tenemos list comprehensions, generadores ni algunas cosas más. Permitidme que resalte una frase sacada de la página principal de numba:

"With a few annotations, array-oriented and math-heavy Python code can be just-in-time compiled to native machine instructions, similar in performance to C, C++ and Fortran". [Énfasis mío]

Siento decepcionar a la audiencia pero numba no acelerará todo el código Python que le echemos: está enfocado a operaciones matemáticas con arrays. Aclarado este punto, vamos a ponernos manos a la obra con un ejemplo aplicado :)

Antes de empezar: instalación

Puedes instalar numba en Windows, OS X y Linux con conda usando este comando:
conda install numba
conda se ocupará de instalar una versión correcta de LLVM, así que no tendrás que compilarla tú mismo. Y ya está.
Ahora viene una opinión personal pero que considero importante: si eres usuario de paquetes científicos y aún no estás utilizando conda (o Anaconda) para gestionarlos, estás en la edad de piedra. Me declaro fanboy absoluto de Continuum Analytics por crear una herramienta de código abierto (conda está en GitHub) que soluciona por fin y de una vez por todas los problemas y frustración que hemos tenido como comunidad desde hace 15 años y que Guido y otros se negaron a atajar. Yo llevo en esto solo desde 2011 pero aún recuerdo lo que es intentar compilar SciPy en Windows. Hazte un favor e instala Miniconda.

Acelerando una función con numba

Voy a tomar directamente el ejemplo que usó Kiko para su artículo sobre Cython y vamos a ver cómo podemos utilizar numba (y un poco de astucia) para acelerar esta función:

"Por ejemplo, imaginemos que tenemos que detectar valores mínimos locales dentro de una malla. Los valores mínimos deberán ser simplemente valores más bajos que los que haya en los 8 nodos de su entorno inmediato. En el siguiente gráfico, el nodo en verde será un nodo con un mínimo y en su entorno son todo valores superiores:

(2, 0) (2, 1) (2, 2)
(1, 0) (1. 1) (1, 2)
(0, 0) (0, 1) (0, 2)

¡Vamos allá!

%install_ext http://raw.github.com/jrjohansson/version_information/master/version_information.py
Installed version_information.py. To use it, type:
  %load_ext version_information
%load_ext version_information
%version_information numpy, numba, cython
Software Version
Python 3.4.3 64bit [GCC 4.4.7 20120313 (Red Hat 4.4.7-1)]
IPython 3.0.0
OS Linux 3.18.6 1 ARCH x86_64 with arch
numpy 1.9.2
numba 0.17.0
cython 0.22
Fri Mar 13 13:44:39 2015 CET
Vamos a empezar por importar los paquetes necesarios e inicializar la semilla del generador de números aleatorios:

import numpy as np
import numba

np.random.seed(0)
Creamos nuestro array de datos:

data = np.random.randn(2000, 2000)
Y voy a copiar descaradamente la función de Kiko:

def busca_min(malla):
    minimosx = []
    minimosy = []
    for i in range(1, malla.shape[1]-1):
        for j in range(1, malla.shape[0]-1):
            if (malla[j, i] < malla[j-1, i-1] and
                malla[j, i] < malla[j-1, i] and
                malla[j, i] < malla[j-1, i+1] and
                malla[j, i] < malla[j, i-1] and
                malla[j, i] < malla[j, i+1] and
                malla[j, i] < malla[j+1, i-1] and
                malla[j, i] < malla[j+1, i] and
                malla[j, i] < malla[j+1, i+1]):
                minimosx.append(i)
                minimosy.append(j)

    return np.array(minimosx), np.array(minimosy)
busca_min(data)
(array([   1,    1,    1, ..., 1998, 1998, 1998]),
 array([   1,    3,   11, ..., 1968, 1977, 1985]))

Paso 1: analizar el código

Lo primero que pensé cuando vi esta función es que no me gustaba nada hacer append a esas dos listas tantas veces. Pero a continuación me pregunté si realmente tendrían tantos elementos... averigüémoslo:

mx, my = busca_min(data)
mx.size / data.size
0.11091025
Tenemos que más de un 10 % de los elementos de la matriz cumplen la condición de ser «mínimos locales», así que no es nada despreciable. Esto en nuestro ejemplo hace un total de más de 400 000 elementos:

mx.size
443641
Ahora la idea de crear dos listas y añadir los elementos uno a uno me gusta todavía menos, así que voy a cambiar de enfoque. Lo que voy a hacer va a ser crear otro array, de la misma forma que nuestros datos, y almacenar un valor True en aquellos elementos que cumplan la condición de mínimo local. De esta forma cumplo también una de las reglas de oro de Software Carpentry: "Always initialize from data".

def busca_min_np(malla):
    minimos = np.zeros_like(malla, dtype=bool)
    for i in range(1, malla.shape[1]-1):
        for j in range(1, malla.shape[0]-1):
            if (malla[j, i] < malla[j-1, i-1] and
                malla[j, i] < malla[j-1, i] and
                malla[j, i] < malla[j-1, i+1] and
                malla[j, i] < malla[j, i-1] and
                malla[j, i] < malla[j, i+1] and
                malla[j, i] < malla[j+1, i-1] and
                malla[j, i] < malla[j+1, i] and
                malla[j, i] < malla[j+1, i+1]):
                minimos[i, j] = True

    return np.nonzero(minimos)
Encima puedo aprovechar la estupenda función nonzero de NumPy. Compruebo que las salidas son iguales:

np.testing.assert_array_equal(busca_min(data)[0], busca_min_np(data)[0])
np.testing.assert_array_equal(busca_min(data)[1], busca_min_np(data)[1])
Y evalúo los rendimientos:

%timeit busca_min(data)
1 loops, best of 3: 4.75 s per loop
%timeit busca_min_np(data)
1 loops, best of 3: 4.62 s per loop
Parece que los tiempos son más o menos parecidos, pero al menos ya no tengo dos objetos en memoria que van a crecer de manera aleatoria. Vamos a ver ahora cómo nos puede ayudar numba a acelerar este código.

Paso 2: aplicando numba.jit(nopython=True)

Como hemos dicho antes, vamos a forzar que numba funcione en modo nopython para garantizar que obtenemos una mejora en el rendimiento. Si intentamos compilar la función definida en primer lugar va a fallar, porque ya hemos dicho más arriba que una de las condiciones es que no se puede asignar memoria a objetos nuevos:

busca_min_jit = numba.jit(nopython=True)(busca_min)
busca_min_jit(data)
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-14-3ca127791a45> in <module>()
      1 busca_min_jit = numba.jit(nopython=True)(busca_min)
----> 2 busca_min_jit(data)

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/dispatcher.py in _compile_for_args(self, *args, **kws)
    155         assert not kws
    156         sig = tuple([self.typeof_pyval(a) for a in args])
--> 157         return self.compile(sig)
    158 
    159     def inspect_types(self, file=None):

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/dispatcher.py in compile(self, sig)
    275                                           self.py_func,
    276                                           args=args, return_type=return_type,
--> 277                                           flags=flags, locals=self.locals)
    278 
    279             # Check typing error if object mode is used

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in compile_extra(typingctx, targetctx, func, args, return_type, flags, locals, library)
    545     pipeline = Pipeline(typingctx, targetctx, library,
    546                         args, return_type, flags, locals)
--> 547     return pipeline.compile_extra(func)
    548 
    549 

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in compile_extra(self, func)
    291                 raise e
    292 
--> 293         return self.compile_bytecode(bc, func_attr=self.func_attr)
    294 
    295     def compile_bytecode(self, bc, lifted=(),

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in compile_bytecode(self, bc, lifted, func_attr)
    299         self.lifted = lifted
    300         self.func_attr = func_attr
--> 301         return self._compile_bytecode()
    302 
    303     def compile_internal(self, bc, func_attr=DEFAULT_FUNCTION_ATTRIBUTES):

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in _compile_bytecode(self)
    532 
    533         pm.finalize()
--> 534         return pm.run(self.status)
    535 
    536 

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in run(self, status)
    189                     # No more fallback pipelines?
    190                     if is_final_pipeline:
--> 191                         raise patched_exception
    192                     # Go to next fallback pipeline
    193                     else:

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in run(self, status)
    181             for stage, stage_name in self.pipeline_stages[pipeline_name]:
    182                 try:
--> 183                     res = stage()
    184                 except _EarlyPipelineCompletion as e:
    185                     return e.result

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in stage_nopython_frontend(self)
    387                 self.args,
    388                 self.return_type,
--> 389                 self.locals)
    390 
    391         with self.fallback_context('Function "%s" has invalid return type'

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in type_inference_stage(typingctx, interp, args, return_type, locals)
    662         infer.seed_type(k, v)
    663 
--> 664     infer.build_constrain()
    665     infer.propagate()
    666     typemap, restype, calltypes = infer.unify()

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/typeinfer.py in build_constrain(self)
    375         for blk in utils.itervalues(self.blocks):
    376             for inst in blk.body:
--> 377                 self.constrain_statement(inst)
    378 
    379     def propagate(self):

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/typeinfer.py in constrain_statement(self, inst)
    480     def constrain_statement(self, inst):
    481         if isinstance(inst, ir.Assign):
--> 482             self.typeof_assign(inst)
    483         elif isinstance(inst, ir.SetItem):
    484             self.typeof_setitem(inst)

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/typeinfer.py in typeof_assign(self, inst)
    514             self.typeof_global(inst, inst.target, value)
    515         elif isinstance(value, ir.Expr):
--> 516             self.typeof_expr(inst, inst.target, value)
    517         else:
    518             raise NotImplementedError(type(value), value)

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/typeinfer.py in typeof_expr(self, inst, target, expr)
    618                                              loc=inst.loc))
    619         else:
--> 620             raise NotImplementedError(type(expr), expr)
    621 
    622     def typeof_call(self, inst, target, call):

NotImplementedError: Failed at nopython (nopython frontend)
(<class 'numba.ir.Expr'>, build_list(items=[]))
En este caso la traza es inservible y especificar los tipos de entrada no va a ayudar. Solo para verificar, vamos a ver qué pasa con el rendimiento si no forzamos el modo estricto:

busca_min_jit_object = numba.jit()(busca_min)
%timeit busca_min_jit_object(data)
1 loops, best of 3: 4.32 s per loop
Pocas ganancias respecto a la función sin compilar. ¿Qué pasa si intentamos lo mismo con la segunda función?

busca_min_np_jit = numba.jit(nopython=True)(busca_min_np)
busca_min_np_jit(data)
---------------------------------------------------------------------------
UntypedAttributeError                     Traceback (most recent call last)
<ipython-input-17-bb9282bb2bfc> in <module>()
      1 busca_min_np_jit = numba.jit(nopython=True)(busca_min_np)
----> 2 busca_min_np_jit(data)

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/dispatcher.py in _compile_for_args(self, *args, **kws)
    155         assert not kws
    156         sig = tuple([self.typeof_pyval(a) for a in args])
--> 157         return self.compile(sig)
    158 
    159     def inspect_types(self, file=None):

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/dispatcher.py in compile(self, sig)
    275                                           self.py_func,
    276                                           args=args, return_type=return_type,
--> 277                                           flags=flags, locals=self.locals)
    278 
    279             # Check typing error if object mode is used

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in compile_extra(typingctx, targetctx, func, args, return_type, flags, locals, library)
    545     pipeline = Pipeline(typingctx, targetctx, library,
    546                         args, return_type, flags, locals)
--> 547     return pipeline.compile_extra(func)
    548 
    549 

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in compile_extra(self, func)
    291                 raise e
    292 
--> 293         return self.compile_bytecode(bc, func_attr=self.func_attr)
    294 
    295     def compile_bytecode(self, bc, lifted=(),

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in compile_bytecode(self, bc, lifted, func_attr)
    299         self.lifted = lifted
    300         self.func_attr = func_attr
--> 301         return self._compile_bytecode()
    302 
    303     def compile_internal(self, bc, func_attr=DEFAULT_FUNCTION_ATTRIBUTES):

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in _compile_bytecode(self)
    532 
    533         pm.finalize()
--> 534         return pm.run(self.status)
    535 
    536 

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in run(self, status)
    189                     # No more fallback pipelines?
    190                     if is_final_pipeline:
--> 191                         raise patched_exception
    192                     # Go to next fallback pipeline
    193                     else:

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in run(self, status)
    181             for stage, stage_name in self.pipeline_stages[pipeline_name]:
    182                 try:
--> 183                     res = stage()
    184                 except _EarlyPipelineCompletion as e:
    185                     return e.result

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in stage_nopython_frontend(self)
    387                 self.args,
    388                 self.return_type,
--> 389                 self.locals)
    390 
    391         with self.fallback_context('Function "%s" has invalid return type'

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/compiler.py in type_inference_stage(typingctx, interp, args, return_type, locals)
    663 
    664     infer.build_constrain()
--> 665     infer.propagate()
    666     typemap, restype, calltypes = infer.unify()
    667 

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/typeinfer.py in propagate(self)
    388                 print("propagate".center(80, '-'))
    389             oldtoken = newtoken
--> 390             self.constrains.propagate(self.context, self.typevars)
    391             newtoken = self.get_state_token()
    392             if config.DEBUG:

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/typeinfer.py in propagate(self, context, typevars)
    110         for constrain in self.constrains:
    111             try:
--> 112                 constrain(context, typevars)
    113             except TypingError:
    114                 raise

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/typeinfer.py in __call__(self, context, typevars)
    267         for ty in valtys:
    268             try:
--> 269                 attrty = context.resolve_getattr(value=ty, attr=self.attr)
    270             except KeyError:
    271                 args = (self.attr, ty, self.value.name, self.inst)

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/typing/context.py in resolve_getattr(self, value, attr)
     82                 raise
     83 
---> 84         ret = attrinfo.resolve(value, attr)
     85         assert ret
     86         return ret

/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numba/typing/templates.py in resolve(self, value, attr)
    241         ret = self._resolve(value, attr)
    242         if ret is None:
--> 243             raise UntypedAttributeError(value=value, attr=attr)
    244         return ret
    245 

UntypedAttributeError: Failed at nopython (nopython frontend)
Unknown attribute "zeros_like" of type Module(<module 'numpy' from '/home/juanlu/.miniconda3/envs/py34/lib/python3.4/site-packages/numpy/__init__.py'>)
Me dice que no conoce la función zeros_like. Si acudimos a la documentación, podemos ver las características de NumPy soportadas por numba y las funciones de creación de arrays no figuran entre ellas. Esto es consistente con lo que hemos dicho más arriba: no vamos a poder asignar memoria a objetos nuevos.

Paso 3: Reestructurar el código

¿Estamos en un callejón sin salida entonces? ¡En absoluto! Lo que vamos a hacer va a ser separar la parte intensiva de la función para aplicar numba.jit sobre ella, e inicializar todos los valores desde fuera. Para los que hayan usado subrutinas en Fortran este enfoque les resultará familiar :)

def busca_min_np_jit(malla):
    minimos = np.zeros_like(malla, dtype=bool)

    _busca_min(malla, minimos)

    return np.nonzero(minimos)

@numba.jit(nopython=True)
def _busca_min(malla, minimos):
    for i in range(1, malla.shape[1]-1):
        for j in range(1, malla.shape[0]-1):
            if (malla[j, i] < malla[j-1, i-1] and
                malla[j, i] < malla[j-1, i] and
                malla[j, i] < malla[j-1, i+1] and
                malla[j, i] < malla[j, i-1] and
                malla[j, i] < malla[j, i+1] and
                malla[j, i] < malla[j+1, i-1] and
                malla[j, i] < malla[j+1, i] and
                malla[j, i] < malla[j+1, i+1]):
                minimos[i, j] = True
Veamos qué ocurre ahora:

busca_min_np_jit(data)
(array([   1,    1,    1, ..., 1998, 1998, 1998]),
 array([   1,    3,   11, ..., 1968, 1977, 1985]))
np.testing.assert_array_equal(busca_min(data)[0], busca_min_np_jit(data)[0])
np.testing.assert_array_equal(busca_min(data)[1], busca_min_np_jit(data)[1])
%timeit busca_min_np_jit(data)
10 loops, best of 3: 62.9 ms per loop
Habéis leído bien: 70x más rápido :)
¡Lo hemos conseguido! Ahora nuestro código funciona en numba sin problemas y encima es endemoniadamente rápido. Para completar la comparación en mi ordenador, voy a reproducir también la función hecha en Cython:

%load_ext cython
%%cython --name probandocython9
import numpy as np
cimport numpy as np
from cpython cimport array as c_array
from array import array
cimport cython

@cython.boundscheck(False) 
@cython.wraparound(False)
cpdef tuple busca_min_cython9(double [:,:] malla):
    cdef c_array.array minimosx, minimosy
    cdef unsigned int i, j
    cdef unsigned int ii = malla.shape[1]-1
    cdef unsigned int jj = malla.shape[0]-1
    cdef unsigned int start = 1
    #cdef float [:, :] malla_view = malla
    minimosx = array('L', [])
    minimosy = array('L', []) 
    for i in range(start, ii):
        for j in range(start, jj):
            if (malla[j, i] < malla[j-1, i-1] and
                malla[j, i] < malla[j-1, i] and
                malla[j, i] < malla[j-1, i+1] and
                malla[j, i] < malla[j, i-1] and
                malla[j, i] < malla[j, i+1] and
                malla[j, i] < malla[j+1, i-1] and
                malla[j, i] < malla[j+1, i] and
                malla[j, i] < malla[j+1, i+1]):
                minimosx.append(i)
                minimosy.append(j)

    return np.array(minimosx), np.array(minimosy)
%timeit busca_min_cython9(data)
10 loops, best of 3: 151 ms per loop
Por tanto, vemos que la versión con numba es el doble de rápida que la versión con Cython. Sobre gustos no hay nada escrito: yo por ejemplo valoro no «salirme» de Python usando numba mientras que a otro puede no importarle incluir especificaciones de tipos como en Cython. Los números, eso sí, son los números :)

Más casos de éxito

La atmósfera estándar

El cálculo de propiedades termodinámicas de la atmósfera estándar es un problema clásico que todo aeronáutico ha afrontado alguna vez muy al principio de su carrera formativa. La teoría es simple: imponemos una ley de variación de la temperatura con la altura \(T = T(h)\), la presión se obtiene por consideraciones hidrostáticas \(p = p(T)\) y la densidad por la ecuación de los gases ideales \(\rho = \rho(p, T)\). La particularidad de la atmósfera estándar es que imponemos que la variación de la temperatura con la altura es una función simplificada y definida a trozos, así que calcular temperatura, presión y densidad dada una altura se parece mucho a hacer esto:

if 0.0 <= h < 11000.0:
    T = T0 + alpha * h
    p = ...  # Algo que depende de T
    rho = p / (R_a * T)
elif 11000.0 <= h < 20000.0:
    T = T1
    p = ...
    rho = p / (R_a * T)
elif 20000.0 <= h <= 32000.0:
    ...
El problema viene cuando se quiere vectorizar esta función y permitir que h pueda ser un array de alturas. Esto es muy conveniente cuando queremos pintar alguna propiedad con matplotlib, por ejemplo.
Se intuye que hay dos formas de hacer esto: utilizando funciones de NumPy o iterando por cada elemento del array. La primera solución se hace farragosa, y la segunda, gracias a la proverbial lentitud de Python, es extremadamente lenta. Mi amigo Álex y yo llevamos pensando sobre este problema años, y nunca hemos llegado a una solución satisfactoria (incluso encontramos algunos bugs en numpy.piecewise por el camino). Este año decidimos cerrar este asunto definitivamente así que con el equipo AeroPython exploramos varias implementaciones distintas. Hasta que por fin lo conseguimos: usamos numba para acelerar los bucles.
numba gana a C++
Como podéis leer en la discusión original, la función de la primera columna está escrita en C++. ¿Impresionado? ;)

Solución de Navier de una placa plana

Para mi proyecto fin de carrera me encontré con la necesidad de calcular la deflexión de una placa rectangular, simplemente apoyada en sus cuatro bordes (es decir, los bordes pueden girar: no están empotrados) sometida a una carga transversal. Este problema tiene solución analítica conocida desde hace tiempo, hallada por Navier:
\(\displaystyle w(x,y) = \sum_{m=1}^\infty \sum_{n=1}^\infty \frac{a_{mn}}{\pi^4 D}\,\left(\frac{m^2}{a^2}+\frac{n^2}{b^2}\right)^{-2}\,\sin\frac{m \pi x}{a}\sin\frac{n \pi y}{b}\)

siendo \(a_{mn}\) los coeficientes de Fourier de la carga aplicada. Como veis, para cada punto \((x, y)\) hay que hacer una doble suma en serie; si encima queremos evaluar esto en un meshgrid, necesitamos un cuádruple bucle. Ya se anticipa que por muy hábiles que estemos, a Python le va a costar.
La clave estuvo, una vez más, en usar numba para optimizar los bucles. En GitHub tenéis el código completo, pero la parte importante es esta:

@numba.jit(nopython=True)
def a_mn_point(P, a, b, xi, eta, mm, nn):
    """Navier series coefficient for concentrated load.

    """
    return 4 * P * sin(mm * pi * xi / a) * sin(nn * pi * eta / b) / (a * b)
 
 
@numba.jit(nopython=True)
def plate_displacement(xx, yy, ww, a, b, P, xi, eta, D, max_m, max_n):
    max_i, max_j = ww.shape
    for mm in range(1, max_m):
        for nn in range(1, max_n):
            for ii in range(max_i):
                for jj in range(max_j):
                    a_mn = a_mn_point(P, a, b, xi, eta, mm, nn)
                    ww[ii, jj] += (a_mn / (mm**2 / a**2 + nn**2 / b**2)**2
                                   * sin(mm * pi * xx[ii, jj] / a)
                                   * sin(nn * pi * yy[ii, jj] / b)
                                   / (pi**4 * D)) 
Solución de una placa plana

Podéis comprobar vosotros mismos que las diferencias de rendimiento en este caso son brutales. Y solo hemos añadido una línea a cada función.

Conclusiones

numba aún no es una herramienta estable, pero está rápidamente alcanzando un grado de madurez suficiente para optimizar código orientado a operar con arrays. Gracias a conda es trivial de instalar y los resultados respecto a soluciones más maduras como Cython son aplastantes, tanto en velocidad de ejecución como en la complejidad del código resultante.
De momento yo me quedo con numba, ¿y tú? ;)

Juan Luis Cano

Estudiante de ingeniería aeronáutica y con pasión por la programación y el software libre. Obsesionado con mejorar los pequeños detalles y con ganas de cambiar el mundo. Divulgando Python en español a través de Pybonacci y la asociación Python España.

More Posts - Website

Follow Me:
TwitterLinkedIn

15 thoughts on “Cómo acelerar tu código Python con numba

  1. Gran artículo!!! Pensar un poco ayuda a resolver los problemas…

    La intención del artículo sobre cython era mostrar sus posibilidades y no está exprimido del todo… (me apunto una actualización para ver si lo puedo optimizar más o aceptar definitivamente el ‘zas, en toda la boca’ con deportividad).

    En este enlace hacen un resumen, creo que acertado, sobre algunas de las posibilidades para ganar rendimiento: http://www.reddit.com/r/Python/comments/2ywf48/cython_numba_f2py_or_other_for_speeding_up/cpdp84q

    Vienen tiempos interesantes!!!

    1. ¡Gracias Kiko! De eso se trata, de ir exprimiendo cada vez más y de hacer *pair programming* :P Seguro que eres capaz de exprimir la versión de Cython todavía más. A numba todavía le queda una cosa que sí tiene Cython: las anotaciones con salida HTML. Estoy esperando impaciente a que lo implementen.

  2. Me sonó a regaño eso de “[...] estás en la edad de piedra. [...] Hazte un favor e instala Miniconda. [...]” pero nunca tan efectivo! :) muchas gracias por eso! Hace algún tiempo quise probar Numba y me fue imposible instalarlo sin fallos. Doblemente útil este artículo, tanto por Numba como por Conda (acabo de instalarlo, y ya instalé además las bibliotecas que uso en mi tesina, a usarlo a partir de mañana, fue realmente instantáneo, y encima ahora tengo Python 3.4 (chau 3.2!)).

    1. Jajaja ¡era el efecto buscado! :D Muchas gracias por el comentario Cristian, lo mismo tienes oportunidad de estudiar este tema con Pybonacci en persona dentro de dos meses… ;)

  3. ¡Interesantísimo artículo! Muy bien descrita la cadena de problemas aparecidos por el camino.Y el aroma a “hay mil cosas por intentar con todas las herramientas de las que disponemos” me gustó porque es de lo que se trata: probar probar probar y probar.Y sobre todo también la sensación que transmites al resolver un problema(y mas si es arduo),que es algo realmente bonito .¡Enhorabuena!

    1. ¡Gracias por el comentario Alberto! Me alegro de que te haya gustado la entrada. Esta ha costado lo suyo pero finalmente encontré la inspiración :P ¡Seguiremos informando, un saludo!

  4. Saludos Juan Luis, esta muy interesante tu articulo, yo he estado intentando acelerar mi codigo Python y he estado probando varias alternativas (jit compilers, cython, numexpr), pero al parecer (en mis resultados), estas solo son significativamente mejores si se trata de cálculos vectoriales o matriciales muy extensos, cuando se programa algo que conlleva cálculos matriciales o vectoriales de forma repetitiva pero al fin y al cabo matrices y vectores pequeños el tiempo de ejecución es ligeramente mayor con respecto a Python original, creo que es algo que hay que acotar, ya que la mayoría de ejemplos que he visto por blogs y demás, es calcular un vector de 12489826481241234231 elementos (entre otros) y compararlo con lo que seria hacerlo con programación en Python pura. Pero nadie habla del limite inferior en el que el tiempo de resolución de ambos empiezan a ser iguales.

    Una desventaja si le pudiera llamar así (en mi punto de vista lo es), es que te tienes que casar con Anaconda o miniconda, ya que tratar de instalarlo por otro lado es un dolor de cabeza y en windows es peor el dolor, en cuanto a scipy, yo no he tenido problemas para “compilarlo”, no se exactamente que significa cuando lo mencionaste al principio, pero ya trae instaladores o paquetes para ambos sistemas operativos (windows y linux) y con py2exe puedes hacer una distribución de tus aplicaciones que corra de forma independiente usando scipy sin ningún problema también de forma independiente.

    Pero en general esta entrada es un buen aporte para aquellas personas que hacen calculos bastantes extensivos, Saludos

    1. Hola Antonio, ¡gracias por el comentario!

      Respecto a lo primero que dices: por un lado, el ejemplo de la atmósfera estándar muestra claramente cómo escalan los tiempos de ejecución (a lo mejor debería poner una gráfica), y puedes ver que numba es la implementación más rápida incluso para argumentos escalares.

      Por otro lado, el ejemplo de la placa de Navier es un cuádruple bucle y las mejoras de rendimiento se notan drásticamente incluso con mallados de una docena de puntos, así que de nuevo no estoy tomando un ejemplo «fabricado» con matrices enormes. Te recomiendo que hagas la prueba tú mismo.

      Respecto a lo que comentas de casarse con Anaconda, al final hay que casarse con alguien. Yo estuve un tiempo casado con Christoph Gohlke y sus paquetes binarios no oficiales, y la situación no me gustaba mucho. Mi matrimonio con pip fue un fiasco también, de Canopy me acabé divorciando para siempre… si no sabes lo que es compilar SciPy, te animo a que descargues el código fuente y hagas `python setup.py install`.

      ¡Un saludo!

      1. y que me dices de easy_install? yo a la vieja escuela me quede con el propio idle original y bueno uso una version un poco mejorada llama idlex, pero no llamaría a esto casarse como tal, e instalo todo con easy_install, cuando hay que hacer alguna compilación en c propiamente de la librería es recomendable instalar Visual C++ Compiler for Python (gratuito), y easy_install se encarga de la compilación (esto para windows claro), y para linux basta con instalar g++, e incluso he podido instalar el ipython 3 con el easy_install en ambos sistemas operativos, la verdad no me ha sido infiel hehehe, pero bueno el punto de la variedad y diversidad es lo que importa, saludos!

  5. Saludos! Quisiera agradecerles por este espacio y el gran aporte que hacen. Hace un mes no sabía nada de Python y ahora he podido trasladar casi todo mis trabajos desde Matlab.
    Mi pregunta es: Es posible combinar Plotly, Numba y los Widgets de iPython? Yo lo he intentado sin éxito.
    Si existe algún tutorial al respecto les agradecería si pudieran compartirlo.

    Muchas gracias!

    1. ¡Hola Javier! Muchas gracias por tus comentarios :)

      La verdad es que no hemos escrito un tutorial tan específico, pero si eres un poco más específico sobre qué problema estás teniendo podemos ayudarte. Te recomiendo que escribas a esta lista de correo:

      python-es@python.org

      Supongo que ya conoces nuestros artículos sobre numba y sobre plotly. Esperamos que tengas éxito, y si te sale bien, que nos mandes un tutorial para publicarlo en el blog :)

      ¡Un saludo!

Leave a Reply