Python for loop numpy vectorization

Question

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)

Answer 1

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