Expression Evaluation (User Defined Functions)#

In the tutorials section, we presented the evaluation of expressions with user defined functions. As no surprise, in this benchmark we are going to compare the performance of ironArray with NumPy and Numba using the same data as in the Expression Evaluation (User Defined Functions) tutorial.

Let’s get down to it!

[1]:
%load_ext memprofiler
%matplotlib inline
import numpy as np
import iarray as ia
import math

ironArray#

[2]:
!vmtouch -e "../tutorials/precip1.iarr" "../tutorials/precip2.iarr" "../tutorials/precip3.iarr"
           Files: 3
     Directories: 0
   Evicted Pages: 215221 (840M)
         Elapsed: 2.4e-05 seconds
[3]:
%%time
precip1 = ia.load("../tutorials/precip1.iarr")
precip2 = ia.load("../tutorials/precip2.iarr")
precip3 = ia.load("../tutorials/precip3.iarr")
CPU times: user 0 ns, sys: 350 ms, total: 350 ms
Wall time: 594 ms

Like in the tutorial we define the same function that computes the mean for the data:

[4]:
from iarray.udf import jit, Array, float32

@jit()
def mean(out: Array(float32, 3),
         p1: Array(float32, 3),
         p2: Array(float32, 3),
         p3: Array(float32, 3)) -> int:

    l = p1.window_shape[0]
    m = p1.window_shape[1]
    n = p1.window_shape[2]

    for i in range(l):
        for j in range(m):
            for k in range(n):
                value = p1[i,j,k] + p2[i,j,k] + p3[i,j,k]
                out[i,j,k] = value / 3

    return 0

and create the ironArray expression from this UDF:

[5]:
%%time
precip_mean_expr = ia.expr_from_udf(mean, [precip1, precip2, precip3])
CPU times: user 11.9 ms, sys: 7.86 ms, total: 19.7 ms
Wall time: 19.5 ms
[6]:
%%mprof_run 1.iarray::mean_UDF
precip_mean = precip_mean_expr.eval()
precip_mean
[6]:
<IArray (720, 721, 1440) np.float32>
memprofiler: used 841.16 MiB RAM (peak of 841.16 MiB) in 0.3453 s, total RAM usage 1930.34 MiB

We will also compute the same UDF but favoring speed and compression ratio.

[7]:
%%time
precip_mean_expr_speed = ia.expr_from_udf(mean, [precip1, precip2, precip3], favor=ia.Favor.SPEED)
CPU times: user 15.8 ms, sys: 7.97 ms, total: 23.8 ms
Wall time: 41 ms
[8]:
%%mprof_run 1.iarray::mean_UDF_speed
precip_mean_speed = precip_mean_expr_speed.eval()
precip_mean_speed
[8]:
<IArray (720, 721, 1440) np.float32>
memprofiler: used 819.42 MiB RAM (peak of 819.42 MiB) in 0.3264 s, total RAM usage 2750.33 MiB
[9]:
%%time
precip_mean_expr_cratio = ia.expr_from_udf(mean, [precip1, precip2, precip3], favor=ia.Favor.CRATIO)
CPU times: user 22.5 ms, sys: 8.52 ms, total: 31.1 ms
Wall time: 51 ms
[10]:
%%mprof_run 1.iarray::mean_UDF_cratio
precip_mean_cratio = precip_mean_expr_cratio.eval()
precip_mean_cratio
[10]:
<IArray (720, 721, 1440) np.float32>
memprofiler: used 700.34 MiB RAM (peak of 700.34 MiB) in 1.8271 s, total RAM usage 3451.20 MiB
[11]:
precip_expr2 = (precip1 + precip2 + precip3) / 3
[12]:
%%mprof_run 1.iarray::mean_lazy
precip_mean2 = precip_expr2.eval()
precip_mean2
[12]:
<IArray (720, 721, 1440) np.float32>
memprofiler: used 859.33 MiB RAM (peak of 859.33 MiB) in 0.3954 s, total RAM usage 4310.54 MiB

NumPy#

Let’s see how NumPy does on this:

[13]:
%%mprof_run
np_precip1 = precip1.data
np_precip2 = precip2.data
np_precip3 = precip3.data
memprofiler: used 8574.96 MiB RAM (peak of 8574.96 MiB) in 5.5940 s, total RAM usage 12885.51 MiB
[14]:
%%mprof_run 2.numpy::mean
np_result = (np_precip1 + np_precip2 + np_precip3) / 3
memprofiler: used 2852.08 MiB RAM (peak of 2852.08 MiB) in 1.0346 s, total RAM usage 15737.64 MiB

Numba#

[15]:
import numba as nb
@nb.jit(nopython=True, parallel=True)
def mean_numba(x, y, z):
    out = np.empty(x.shape, x.dtype)
    for i in nb.prange(x.shape[0]):
        for j in nb.prange(x.shape[1]):
            for k in nb.prange(x.shape[2]):
                out[i, j, k] = (x[i, j, k] + y[i, j, k] + z[i, j, k]) / 3
    return out

In this example we are enforcing numba to execute outside the Python interpreter (nopython=True) for leveraging the parallelism support in Numba (parallel=True). We will compute the mean twice because the first time Numba loses some time initializing buffers.

[16]:
%%mprof_run
nb_mean = mean_numba(np_precip1, np_precip2, np_precip3)
memprofiler: used 2864.84 MiB RAM (peak of 2864.84 MiB) in 0.8914 s, total RAM usage 18622.80 MiB
[17]:
# Avoid tracking memory consumption of this object
del nb_mean
[18]:
%%mprof_run 3.numba::mean
nb_mean = mean_numba(np_precip1, np_precip2, np_precip3)
memprofiler: used 2849.71 MiB RAM (peak of 2849.71 MiB) in 0.5389 s, total RAM usage 18623.16 MiB

Results#

[19]:
%mprof_plot .*::mean -t "Mean computation"