Why Python code takes so long to find close points in a dataset?

Question

I am trying to make an efficient Python code that, given a set of data points in the form

a1 a2 a3 ... an (point 1)
b1 b2 b3 ... bn (point 2)
         .
         .
         .

will find which of these are too close (within a threshold value).

I already have the following Fortran routine that solves this problem:

program check_closeness                                           
    implicit none
    integer, parameter :: npoints = 10000, ndis = 20
    real(8), parameter :: r = 1e0, maxv = 3e0, minv = 0e0
    real(8), dimension(npoints,ndis) :: dis
    real(8) :: start, ende

    call RANDOM_NUMBER(dis)
    dis = (maxv - minv) * dis + minv

    call cpu_time(start)
    call remove_close(npoints, ndis,dis, r)
    call cpu_time(ende)
    write(*,*) 'Time elapsed', ende - start
endprogram

subroutine remove_close(npoints, ndis, points, r)
    implicit none
    integer, intent(in) :: npoints, ndis
    integer :: i, i_check
    real(8), dimension(npoints, ndis), intent(in) :: points
    real(8), intent(in) :: r
    logical :: is_close

    do i=1,npoints
       is_close = .FALSE.
       i_check = i

       do while (.not. is_close .and. i_check < npoints)
          i_check = i_check + 1
          is_close = all(abs(points(i,:) - points(i_check,:)) < r)
       enddo
    enddo
end subroutine

This execution takes in my computer: Time elapsed 1.0651889999999999 (seconds).

Now, I write the same exact code (or so I think) in Python:

 import numpy as np                                       
 import time
 
 def check_close(a, r):
     npoints = a.shape[0]
     for i, vec1 in enumerate(a):
         is_close = False
         icheck = i
 
         while (not is_close and icheck < npoints - 1):
             icheck +=1
             vec2 = a[icheck,:]
             is_close = all(np.abs(vec1 - vec2) < r)
 
 maxv, minv = 3, 0
 a = np.random.rand(10000, 20)
 a = (maxv - minv) * a + minv
 r = 1e0
 
 ini = time.time()
 check_close(a, r)
 fin = time.time()
 print('Time elapsed {}'.format(fin - ini))

This execution takes Time elapsed 102.60617995262146 (seconds), which is considerably slower than the Fortran one.

I tried yet another Python version of this routine which is considerably faster, but still does not approach the Fortran one:

import numpy as np                        
import time

def check_close(a, r):
    for i, vec1 in enumerate(a[:-1]):
        d = np.abs(a[i+1:,:] - vec1)
        is_close = np.any(np.all(d < r, axis=1))

maxv, minv = 3, 0
a = np.random.rand(10000, 20)
a = (maxv - minv) * a + minv
r = 1e0

ini = time.time()
check_close(a, r)
fin = time.time()
print('Time elapsed {}'.format(fin - ini))

In this case the execution takes Time elapsed 3.4987785816192627 (seconds). From this, I guess the improvement comes from the vectorization achieved and removing the while loop. On the other hand, this implementation does not benefit from stopping the search as one close point is found.

My questions are:

1. What is it that takes so much time in the Python implementations?
2. Is there any way I could rewrite the Python code to make if (almost) as fast as the Fortran one?

[Edit]

I was asked to use time.perf_counter() to measure executions times in Python codes.

Now I use:

ini = time.perf_counter()
check_close(a, r)
fin = time.perf_counter()
print('Time elapsed {}'.format(fin - ini))

to measure them and the updated times are:

Python implementation 1: Time elapsed 98.63728801719844 (seconds)

Python implementation 2: Time elapsed 3.3923211600631475 (seconds)

Answer 1

Numpy was created specifically BECAUSE python is incredibly slow. As you showed, fortran code is much faster without much algorithmic differences. So to answer your questions:

1. Python is a scripting language and was not built around speed. If you need speed, use libraries such as numpy, like you are doing, or try to incorporate native C/C++ code into your program where you need it.
2. To make it as fast as fortran, try to find a way to only use numpy routines, getting rid of the for loop. This will drastically increase your performance.

Try to think about your problems and how you can solve them without loops; this will be the key to unlock speed in python.

By the way, I am not sure what you intend with your use of np.abs() but it is not calculating the euclidean distance between two points. Use np.linalg.norm() for that.

Answer 2

You can do better in Python, but it takes some work. Surprisingly I think your approach using only NumPy is about as good as you can get without involving other modules. However, if you are willing to use numba you can do better than even Fortran.

The function below returns the close vectors. It is effectively C-code written in Python that is then compiled to native using numba.

import numpy as np
from math import fabs
from numba import njit

@njit
def check_close_nb2(a, r):
    n = a.shape[0]
    m = a.shape[1]
    # Selection array if you want to use it later
    close = np.zeros(n, dtype=np.bool_)
    for i in range(n):
        for j in range(i + 1, n):
            for k in range(m):
                # Short circuit comparison
                last = fabs(a[i, k] - a[j, k]) >= r
                if last:
                    break
                # If we reach here we found a close
            if k == (m - 1) and not last:
                # Once a close value is found stop looking
                close[i] = 1
                break
    return close

When I run this on my machine using %timeit I get

%timeit check_close_nb2(a, r)
544 ms ± 4.07 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

When I run your code above, I see

%timeit check_close(a, r)
4.91 s ± 69.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

so my comparison suggests that my machine is a bit slower than yours, and would suggest a time around 400ms for the numba version.

Answer 3

I have tried 2 methods, both under 3 seconds on my machine, but not yet as fast as FORTRAN.

I generate similar random data, but increased the r as otherwise I got zero cases

import time

import numpy as np

maxv, minv = 3, 0

a = np.random.rand(10000, 20)
a = (maxv - minv) * a + minv
r = 3e0

1 - BallTree (sklearn)

from sklearn.neighbors import BallTree

ini = time.time()
tree = BallTree(a, leaf_size=5000)

within = tree.query_radius(a, r=r)

fin = time.time()

print('Balltree (sklearn) based - time elapsed {}'.format(fin - ini))

Which gave

Balltree (sklearn) based - time elapsed 1.5845096111297607

edit: Weirdly (can someone explain?) an insane high leafsize=5000 worked best in my case. Normally the leafsize is around 10.

2 - pdist (scipy)

from scipy.spatial.distance import pdist


ini = time.time()

condensed_distance_small = pdist(a) <= r
condensed_indici = np.argwhere(condensed_distance_small)

n = 10000

#
# i,j contain the indici of points that are within r of each other
# based on condensed-indici formulas of
# https://stackoverflow.com/questions/13079563/how-does-condensed-distance-matrix-work-pdist

i = np.ceil((1/2.) * (- (-8*condensed_indici + 4 *n**2 -4*n - 7)**0.5 + 2*n -1) - 1).astype(int)
elem_in_i_rows = (i+1) * (n - 1 - (i+1)) + ((i+ 1)*(i + 2))//2
j = (n - elem_in_i_rows + condensed_indici)

fin = time.time()

print('pdist (scipy) based - time elapsed {}'.format(fin - ini))

Which gave

pdist (scipy) based - time elapsed 0.6769788265228271

the pdist is returning condensed form indici of point, so there is some (i,j) index administration to be done, but this is done with vectors=fast.