pythonnumpyvectorization

Python for loop numpy vectorization


Do you have any suggestions how I can speed up the "f_img_update" function in the code below? I tried adding numba @jit decorator, but without luck. I guess that numpy vectorization might work, but I am not enough educated to make it work on my own :/

Code:

import numpy as np
import time
import random
from random import randrange


dimensions = (100, 50)  # [m]
resolution = 5  # [cm] 1,2,4,5,10,20,25,50,100
pix_res = int(100 / resolution)

height = dimensions[0] * pix_res
width = dimensions[1] * pix_res
img_size = (height, width)

# Make "RGBA" image
img = np.zeros((*img_size, 4), np.uint16)

# DATA PREPARATION
xmin = 1700000
ymin = 1400000
zmin = 100

xmax = xmin + dimensions[0]
ymax = ymin + dimensions[1]
zmax = 150

data = []
for i in range(100000):
    xyzi = [random.uniform(xmin, xmax), random.uniform(ymin, ymax), random.uniform(zmin, zmax), randrange(255)]
    data.append(xyzi)

# IMAGE UPDATE
def f_img_update(f_img, f_xmin, f_ymin, f_data, f_pix_per_1m):  # IMAGE UPDATE
    for f_xyzi in f_data:
        f_x_img, f_y_img = f_xyzi[0:2]
        f_x_img = int((f_x_img - f_xmin) * 100 / f_pix_per_1m)
        f_y_img = int((f_y_img - f_ymin) * 100 / f_pix_per_1m)
        f_img[f_x_img][f_y_img] = f_xyzi
    return f_img


t1 = time.time()
img = f_img_update(img, xmin, ymin, data, pix_res)
t2 = time.time()
print("image update time: ", t2 - t1)

Solution

  • The easiest way is to convert the data from standard python list to np.array, then use @njit decorator:

    import random
    import time
    from random import randrange
    
    import numpy as np
    from numba import njit
    
    dimensions = (100, 50)  # [m]
    resolution = 5  # [cm] 1,2,4,5,10,20,25,50,100
    pix_res = int(100 / resolution)
    
    height = dimensions[0] * pix_res
    width = dimensions[1] * pix_res
    img_size = (height, width)
    
    # Make "RGBA" image
    img = np.zeros((*img_size, 4), np.uint16)
    
    # DATA PREPARATION
    xmin = 1700000
    ymin = 1400000
    zmin = 100
    
    xmax = xmin + dimensions[0]
    ymax = ymin + dimensions[1]
    zmax = 150
    
    data = []
    for i in range(100_000):
        xyzi = [
            random.uniform(xmin, xmax),
            random.uniform(ymin, ymax),
            random.uniform(zmin, zmax),
            randrange(255),
        ]
        data.append(xyzi)
    
    data = np.array(data)
    
    
    # IMAGE UPDATE
    def f_img_update(f_img, f_xmin, f_ymin, f_data, f_pix_per_1m):  # IMAGE UPDATE
        for f_xyzi in f_data:
            f_x_img = f_xyzi[0]
            f_y_img = f_xyzi[1]
            f_x_img = int((f_x_img - f_xmin) * 100 / f_pix_per_1m)
            f_y_img = int((f_y_img - f_ymin) * 100 / f_pix_per_1m)
            f_img[f_x_img, f_y_img] = f_xyzi
        return f_img
    
    
    @njit
    def f_img_update_numba(f_img, f_xmin, f_ymin, f_data, f_pix_per_1m):  # IMAGE UPDATE
        for f_xyzi in f_data:
            f_x_img = f_xyzi[0]
            f_y_img = f_xyzi[1]
            f_x_img = int((f_x_img - f_xmin) * 100 / f_pix_per_1m)
            f_y_img = int((f_y_img - f_ymin) * 100 / f_pix_per_1m)
            f_img[f_x_img, f_y_img] = f_xyzi
        return f_img
    
    
    # compile it first
    assert np.allclose(
        f_img_update(img, xmin, ymin, data, pix_res),
        f_img_update_numba(img, xmin, ymin, data, pix_res),
    )
    
    img_copy = img.copy()
    t1 = time.time()
    _ = f_img_update(img_copy, xmin, ymin, data, pix_res)
    t2 = time.time()
    print("image update normal: ", t2 - t1)
    
    img_copy = img.copy()
    t1 = time.time()
    _ = f_img_update_numba(img_copy, xmin, ymin, data, pix_res)
    t2 = time.time()
    print("image update time numba: ", t2 - t1)
    

    Prints on my AMD 5700x:

    image update normal:  0.07558226585388184
    image update time numba:  0.001255035400390625