Aplicación de Numba para agilizar el cálculo de irradiancias en la caracterización de lentes ópticas difractivas
Introducción
Los lenguajes de programación que se utilizan para el análisis de fenómenos físicos que requieren una cantidad notable de recursos, son optimizados para aprovechar al máximo la gestión entre el CPU, el GPU y la memoria y realizar las tareas de cálculo en el menor tiempo posible. Python, un lenguaje que ha tomado auge en los últimos años, tiene la notable desventaja respecto de C que requiere, relativamente, mucho tiempo para ejecutar sus instrucciones, ya que es un lenguaje interpretado; es decir, a medida se van leyendo las instrucciones, éstas se van ejecutando.
Por tanto, se han invertido múltiples esfuerzos por diferentes programadores para agilizar la ejecución de programas: Desarrollando un archivo compilado, utilizando la GPU del sistema y aprovechando el manejo de estructuras vectoriales; sin embargo, algunos programas ya se encuentran elaborados por científicos y desarrolladores, lo cual hace más difícil implementar las medidas de ejecución antes mencionadas.
Una opción viable para aumentar notablemente la rapidez de los programas escritos en Python, en la aplicación del paquete numpy, es el uso de una libraría conocida como numba. Esta librería, mediante el uso de unos “decoradores” hace posible que la función o método de una clase, sea analizada y procesada de la manera más optima posible, incluso paralelizando procesos que se encuentran contenidos en una estructura de la forma for‑loop.
El
contexto en el cual se desea agilizar el proceso de cálculo en en la
caracterización de lentes difractivas, donde cada lente está
representada por matrices de
elementos, donde
puede
tener valores entre 200 y 1000. Para el análisis, todos esos
elementos se calculan en ecuaciones como la siguiente
Donde las variables
refieren
a las coordenadas de cada elemento matricial de la lente, la cual debe integrarse y calcularse para cada coordenada del espacio
además,
se incluye la variable
que
especifica a la longitud de onda de la luz aplicada a la lente. La
fórmula anterior, a su vez, se aplica un número de veces igual que
el de puntos en el espacio en que se requiere conocer la irradiancia,
generándose de esta manera planos de irradiancia y árboles de
irradiancia.
Se hace notar que la computadora utilizada para esta ilustración, no posee tarjetas de vídeo que contengan “cuda cores” y ha sido necesario buscar otras alternativas para acelerar el código.
De manera general, para realizar el cálculo de irradiancias en el espacio 3D utilizando Python, se sigue la estructura:
import librerías_utilizadas
…
def function1(**args) # Esta es una función paralelizable
…
def function0(**args) # Esta función organiza la información y asigna valores a variables, llama function1.
…
Se muestran los resultados obtenidos para tres casos. Se hace notar que las funciones representan métodos de una clase llamada lente.
a) Realizando ejecuciones paralelas de function1 mediante la librería de multiprocessing
La función del cálculo de irradiancia se declara mediante function0, una variable tipo String, posteriormente se utiliza la función “pool.apply” para ejecutar el método parallel_xyz de forma paralela. En esta última función, se utiliza eval y un dictionario con las variables que requerirá function0.
import multiprocessing as mp
...
def parallel_xyz(self, axis, function0, dict0):
(dict0['x'], dict0['y'], dict0['z']) = axis
result = eval(function0, None, dict0)
return result
…
function0 ='np.abs(np.exp(1j*k*z)/(1j*lambdaa*z)*np.sum(lente * np.exp(1j*k*((x-x0)**2+((y-y0)**2))/(2*z))))'
pool = mp.Pool()
results = [pool.apply(self.parallel_xyz, args=(u, function0, dict0)) for u in axis]
pool.close()
b) Utilizando el paquete numba sin paralelización.
La función de la irraciancia no se declara como un String como en el caso anterior, ya que numba no reconoce su contenido y reporta un error. Nótese el uso de los decoradores @staticmethod que debe utilizarse porque se aplica al método de una clase y @jit que buscará agilizar el proceso; sin embargo lo realiza con un solo procesador.
Nótese que cada valor asignado a result[idx] se calcula utilizando operaciones entre matrices (lente) y escalares (k)
from numba import jit, jitclass, float64
from numba.errors import NumbaDeprecationWarning, NumbaPendingDeprecationWarning
import warnings
…
@staticmethod
@jit(nopython=True)
def parallel_xyz2(lente, x0, y0, axis, k, lambdaa):
result = np.empty((len(axis)))
for idx,(x, y, z) in enumerate(axis):
result[idx] = np.abs(np.exp(1j*k*z)/(1j*lambdaa*z)*np.sum(lente * np.exp(1j*k*((x-x0)**2+((y-y0)**2))/(2*z))))
return result
…
results = self.parallel_xyz2(self.data, self.dim_x, self.dim_y, axis, self.k, self.lambdaa)
c) Utilizando numba con paralelización.
De manera similar al anterior, se utilizan decoradores, ahora con la particularidad de proporcionar el argumento parallel = True en el decorador @jit. De manera especial, para que las operaciones se desarrollaran paralelamente, ha sido necesario declarar las operaciones de cada i-ésimo elemento
from numba import njit
…
@staticmethod
@njit(parallel=True)
def parallel_xyz2(lente, x0, y0, x, y, z, k, lambdaa):
var = x.size
result = [0.0] * var
for i in prange(var):
result[i] = np.abs(np.exp(1j * k * z[i]) / (1j * lambdaa * z[i]) * np.sum(lente * np.exp(1j * k * ((x[i] - x0) ** 2\ + ((y[i] - y0) ** 2)) / (2 * z[i]))))
return result
…
results = self.parallel_xyz2(self.data, self.dim_x, self.dim_y, x1, y1, uz1, self.k, self.lambdaa)
Para realizar las pruebas, se utilizaron dos
lentes iguales y se calculó la irradiancia en un plano
arbitrario,
el primero de
puntos
y el segundo
puntos
mediante las líneas de código siguientes:
(código 1) irr0 = lente0.irradiancia_xyuz(np.linspace(-0.1, 0.1, 50), np.linspace(-0.1, 0.1, 50), 0.355)
(código 2) irr0 = lente0.irradiancia_xyuz(np.linspace(-0.1, 0.1, 100), np.linspace(-0.1, 0.1, 100), 0.355)
Resultados
Los resultados obtenidos son:
|
a) Multiprocessing [s] |
b) Numba sin paralelización [s] |
c) Numba con paralelización [s] |
Código 1 |
36.506 |
8.324 |
2.988 |
Código 2 |
144.578 |
30.276 |
7.097 |
Si se define la relación de tiempo
puede observarse una relación de
para el código 1 y
para el código 2.
Conclusiones
No siempre se dispone de procesadores gráficos que contengan “cuda cores” de Nvidia para acelerar el código de programas para cálculo de matrices y deben aprovecharse los recursos de Intel o AMD disponibles en un momento determinado.
El uso del paquete numba ha demostrado que acelera notablemente, entre 12.5 y 20 veces, los tiempos requeridos para el cálculo de irradiancias en un plano.
Se recomienda el uso de numba en programas escritos en Python para cálculo de matrices, por su notable nivel de aceleración del código y muy poca necesidad de cambiar los programas previamente escritos.
