How can I speed up a det distance function in python with Numba?-CodePudding

My 'arr_coordinates' stores the x and y coordinates in 2 different columns. The 'distance' function returns an array that contains all distances between the coordinates.

My code:

xlim = (0, 1800)
ylim = (0, 1800)
arr_x = ([])
arr_y = ([])
for i in range(300):
    arr_x = np.append(arr_x, random.randint(xlim[0], xlim[1]))
    arr_y = np.append(arr_y, random.randint(ylim[0], ylim[1]))
arr_coordinates = np.vstack((arr_x, arr_y)).T

@numba.jit(forceobj=True)
def distance(arr_coordinates):
    arr_distances = ([])
    for i in range(len(arr_coordinates)):
        coordinate = arr_coordinates[i]
        for j in range(len(arr_coordinates)):
            other_coordinate = arr_coordinates[j]
            distance = ((other_coordinate[0] - coordinate[0]) ** 2   (other_coordinate[1] - coordinate[1]) ** 2) ** 0.5 #√[(x₂ - x₁)²   (y₂ - y₁)²]
            arr_distances = np.append(arr_distances, distance)
    return arr_distances

print("time", timeit.timeit(functools.partial(distance, arr_coordinates), number=1))

Without the Numba decorator it is faster.

CodePudding user response：

first of all, you should forget a function with the name np.append exists, this function is very slow, it's a thousand times slower than list.append, if it is possible to calculate the output size beforehand and reserve it at the beginning then you should definitely do it, which is currently your case.

import numpy as np
import random
import numba
import timeit
import functools
from math import sqrt

xlim = (0, 1800)
ylim = (0, 1800)
arr_x = ([])
arr_y = ([])
for i in range(300):
    arr_x = np.append(arr_x, random.randint(xlim[0], xlim[1]))
    arr_y = np.append(arr_y, random.randint(ylim[0], ylim[1]))
arr_coordinates = np.vstack((arr_x, arr_y)).T

@numba.jit(forceobj=True)
def distance(arr_coordinates):
    arr_distances = ([])
    for i in range(len(arr_coordinates)):
        coordinate = arr_coordinates[i]
        for j in range(len(arr_coordinates)):
            other_coordinate = arr_coordinates[j]
            distance = ((other_coordinate[0] - coordinate[0]) ** 2   (other_coordinate[1] - coordinate[1]) ** 2) ** 0.5 #√[(x₂ - x₁)²   (y₂ - y₁)²]
            arr_distances = np.append(arr_distances, distance)
    return arr_distances

@numba.njit
def distance_numba(arr_coordinates):
    arr_distances = np.empty((len(arr_coordinates),len(arr_coordinates)),dtype=arr_coordinates.dtype)
    for i in range(len(arr_coordinates)):
        coordinate = arr_coordinates[i]
        for j in range(len(arr_coordinates)):
            other_coordinate = arr_coordinates[j]
            distance = sqrt(((other_coordinate[0] - coordinate[0]) ** 2   (other_coordinate[1] - coordinate[1]) ** 2)) #√[(x₂ - x₁)²   (y₂ - y₁)²]
            arr_distances[i,j] = distance
    return arr_distances

functools.partial(distance, arr_coordinates)()  # never time the first call
functools.partial(distance_numba, arr_coordinates)()  # never time the first call
print("time", timeit.timeit(functools.partial(distance, arr_coordinates),number = 1))
print("time", timeit.timeit(functools.partial(distance_numba, arr_coordinates), number= 1))

time 2.0872249000000003
time 0.0002234000000003178

and there is a 10,000 speedup just because we no longer use np.append, and actually reserving the memory beforehand.

a small note, using sqrt is also faster than **0.5 because computer has special hardware that does sqrt faster than the hardware that does pow.

as for **2 the compiler will correctly resolve it to x*x instead of calling pow so it doesn't need to be hand-unrolled.