scipy_vs_fastdist.py

import matplotlib.pyplot as plt

from scipy.spatial import distance
from fastdist import fastdist
import timeit
import numpy as np
import numba
from fastdist import __version__ as fastdist_version
import scipy
from cpuinfo import get_cpu_info

def dist_scipy(a, b):
    return distance.cdist(a,b,"sqeuclidean")

def dist_fast(a, b):
    return fastdist.matrix_to_matrix_distance(a, b, fastdist.sqeuclidean, "sqeuclidean")


nn = [10,100,1000,10000,100000]
d = 100

t_scipy = []
t_fast = []

centers = np.random.rand(200, d)

for n in nn:
    print(n)
    X = np.random.rand(n,d)
    
    t_scipy.append(
        timeit.repeat(
            "dist(X, centers)",
            number=1,
            repeat=7,
            globals=dict(dist=dist_scipy, X=X, centers=centers)
        )
    )
    
    t_fast.append(
        timeit.repeat(
            "dist(X, centers)",
            number=1,
            repeat=7,
            globals=dict(dist=dist_fast, X=X, centers=centers)
        )
    )
    

nn = np.array(nn)
t_scipy = np.array(t_scipy).min(axis=1)
t_fast = np.array(t_fast).min(axis=1)

fig, ax = plt.subplots()

ax.plot(nn, t_scipy, "-x", label="scipy.spatial.distance.cdist")
ax.plot(nn, t_fast, "-x", label="fastdist.matrix_to_matrix_distance")

ax.set_xscale("log")
ax.set_yscale("log")

cpu_info = get_cpu_info()

info = (
    f"numba version: {numba.__version__}",
    f"fastsdist version: {fastdist_version}",
    f"scipy version: {scipy.__version__}",
    f"CPU: {cpu_info['brand_raw']}",
    "",
    f"dimensions: {d}"
)

ax.text(0.95, 0.05, "\n".join(info), transform=ax.transAxes, fontsize=10,
        verticalalignment='bottom', ha="right",)

ax.set_xlabel("Number of samples")
ax.set_ylabel("Time [s]")

plt.legend()