Dr. Numpy-Love or: How I Learned to Stop Worrying and SIMD in Python
By: Debanjan Basu
Tags: Simd, Numpy, Data Science, Parallel Processing, Python

Problem: Finding the nearest locations on the earth’s surface

We create an array of latitudes ( -90 < lat < 90 ) and longitudes ( -180 < lng < 180 ) . The task is to find the nearest locations corresponding to each one and reporting the index and distance between them.

Solution — Step 1. Create the array of locations

import numpy as np

import math np.random.seed(42)

N = 1000 def generate_random_locations(N):

"""Generates a numpy array of `N` locations (:= lat, lng) as 2-tuples"""

locations = np.random.random((N, 2)) - 0.5

return np.matmul(locations, np.diag([180, 360])) locations_100 = generate_random_locations(N)

locations_500 = generate_random_locations(5*N)

We created two datasets here — I went with a 100 locations for quickly evaluating the accuracy of the answers and a larger dataset of 500 locations to benchmark the execution times.

Solution — Step 2. Use a pure python solution for reference

This is the kind of unit-testing that really works for most data science contexts. Having a more tractable but slower solution on smaller datasets not only helps future-proof our algorithms, it also helps us tinker around and find the correct combination of operators that do what we want. I, personally, also like to think of the faster solutions as X-fold performance improvements on the slower but more readable code — this makes it really exciting for me.

Then we compute the distances between each pair of points and accumulate the nearest locations for each row.

Expression for Haversine distance in spherical polar coordinates. θ≡ lat in radians,ϕ≡ lng in radians

def haversine_distance(lat1, lng1,lat2, lng2, r=6378.1370):

"""Returns the haversine distance between two points."""



# convert from degrees to radians

lat1, lng1, lat2, lng2 = [

math.radians(coord)

for coord in (lat1, lng1, lat2, lng2)

]

return 2.0*r*math.asin(

math.sqrt(

math.sin((lng2-lng1)/2.0)**2

+ math.cos(lng1)*math.cos(lng2)*math.sin((lat2-lat1)/2.0)**2

)

) def nearest_neighbours(locations, dist=haversine_distance):

"""Returns the index and distance of the nearest neighbours"""

N = len(locations)

results_idx = []

results_dist = []

for iloc1, loc1 in enumerate(locations):

min_dist = None

for iloc2, loc2 in enumerate(locations):

if iloc2 == iloc1:

continue

dist_1_2 = haversine_distance(*loc1, *loc2)

if (min_dist is None) or (min_dist > dist_1_2):

min_dist = dist_1_2

min_idx = iloc2

results_idx.append(min_idx)

results_dist.append(min_dist)

return results_idx, results_dist

Notice that all pairs of locations need to be checked exactly once. Since we are not saving the distances to be reused later, this algorithm has no space-complexity at all.

Solution — Step 3a. Use a numpy solution with vectorized distance metric for reference

def nearest_neighbours_vectorized(locations, dist):

"""Returns the index and distance of the nearest neighbours"""

N = len(locations)

lat_mat = np.tile(locations[:, 0], [N, 1],)

lng_mat = np.tile(locations[:, 1], [N, 1],)

distance_matrix = dist(lat_mat, lng_mat, lat_mat.T, lng_mat.T)

np.fill_diagonal(distance_matrix, np.nan)

results_idx = np.nanargmin(distance_matrix, axis=0)

results_dist = distance_matrix[results_idx, np.arange(N)]

return results_idx, results_dist

Notice the expression distance_matrix = dist(lat_mat, lng_mat, lat_mat.T, lng_mat.T) … we need to note that —

Distance Matrix is symmetric: We need to recalculate the distances between the pairs (i, j) and (j, i) . As far as I know, matrices with symmetries (not just symmetric matrices, but other kind of symmetries that are preserved under multiplication or other operations) are not really supported in NumPy. But we could have reduced the size of the problem to a finite factor if that were possible. In general however, the improvement from symmetries are nothing comparable to the raw power that SIMD unleashes in such cases. Broadcasting: Instead of calculating the D_ij separately over a loop, we have restructured and passed in the relevant locations and the transpose of locations using np.tile (also possible with np.vstack) . This is still overcounting the (i, j) and (j, i) pairs, but numpy can’t do better AFAIK.

Solution — Step 3b. Define a custom numpy-only vectorized haversine distance function

def haversine_distance_vectorized(lat1, lng1,lat2, lng2, r=6378.1370):

# convert from degrees to radians

lat1, lng1, lat2, lng2 = [

np.radians(coord)

for coord in (lat1, lng1, lat2, lng2)

]

return 2.0*r*np.arcsin(

np.sqrt(

np.sin((lng2-lng1)/2.0)**2

+ np.cos(lng1)*np.cos(lng2)*np.sin((lat2-lat1)/2.0)**2

)

)

We implement a numpy-only vectorized function that can handle arrays as input — the interface is same as np.vectorize(haversine_distance) .

Solution — Step 3c. Judge the accuracy visually

Now once I am sure that my program is accurate, I would make this into an unit-test. But for now while developing it, since an yes-no answer is not as instructive as looking at a subset of the output — which has so many more clues!

nearest_locations = nearest_neighbours(locations_100)

list(zip(*nearest_locations))[:10]

# [(69, 870.4978432803006),

# (355, 479.04535249082073),

# (294, 357.76547007124174),

# (264, 62.528270176824584),

# (699, 99.49969360321941),

# (235, 165.90804855573805),

# (935, 82.47250221260545),

# (433, 226.5015940626329),

# (23, 239.28265676064407),

# (134, 392.967205909576)] nearest_locations_vectorized = nearest_neighbours_vectorized(locations_100, haversine_distance_vectorized) list(zip(*nearest_locations_vectorized))[:10] # [(69, 870.4978432803006),

# (355, 479.04535249082073),

# (294, 357.76547007124174),

# (264, 62.528270176824584),

# (699, 99.49969360321941),

# (235, 165.90804855573805),

# (935, 82.47250221260545),

# (433, 226.5015940626329),

# (23, 239.28265676064407),

# (134, 392.967205909576)]

Solution — Step 3d. Let it rip and time-it

%%time

_ = nearest_neighbours(locations_500) # CPU times: user 2min, sys: 279 ms, total: 2min Wall time: 2min %%time

_ = nearest_neighbours_vectorized(locations_500,

np.vectorize(haversine_distance)) # CPU times: user 1min 3s, sys: 1.93 s, total: 1min 5s Wall time: 1min 5s %%time

_ = nearest_neighbours_vectorized(locations_500,

haversine_distance_vectorized) # CPU times: user 3.89 s, sys: 93.4 ms, total: 3.98 s Wall time: 4.01 s

This is where SIMD shines — a speedup of 100% with the vectorized function that still runs python code and a speedup of more than 4000% for a native vectorized function that runs the BLAS libraries in numpy.

Beyond numpy: Scaling up

The immediate tool to trivially parallelize this across clusters would be dask. I would also want to use the 2-fold symmetry for each matrix — for a multiplication of M matrices each with 2-fold symmetry, this reduces the number of operations 2^M fold. This kind of utilization of symmetries could still utilize SIMD efficiently, but I haven’t come across a framework that can do this.

Another interesting speedup would be to use triangular inequality to set up a lower bound for the distances. This way only those pairs of points that clear the lower bound may be considered for the distance caculations. I think this might still be relevant for SIMD if each worker calculates its own part of the grid first (i.e the block-diagonals of the distance_matrix) and then the off-diagonals may be filtered down using triangular inequality.