SIMD Manhattan Distance

benchmark.py

from sklearn.metrics._dist_metrics import DistanceMetric, DistanceMetric32
from statistics import mean, stdev
from time import perf_counter
from functools import partial
from itertools import product
from pathlib import Path
import numpy as np
import csv

results_path = 'local_artifacts/benchmarks/native/quick'
Path(results_path).mkdir(parents=True, exist_ok=True)
branch = "main"

def _generate_PWD_data(n_samples_X, n_samples_Y, n_features, n_classes, n_outs=1, random_state=0):
    rng = np.random.RandomState(random_state)
    X = rng.randn(n_samples_X, n_features)
    Y = rng.randn(n_samples_Y, n_features)
    y_shape = (n_samples_Y,) if n_outs == 1 else (n_samples_Y, n_outs)
    y = rng.randint(n_classes, size=y_shape)

    return X, Y, y

benchmark_config = [
    (
        partial(_generate_PWD_data, n_features=100, n_classes=2),
        product(
            [5_000, 10_000, 20_000],
            [np.float32, np.float64],
        ),
    ),
]


N_REPEATS = 30

with open(f'{results_path}{branch}.csv', 'w', newline='') as csvfile:

    writer = csv.DictWriter(
        csvfile,
        fieldnames=[
            "n_samples",
            "dtype",
            "n_repeat",
            "duration",
        ],
    )
    writer.writeheader()

    for make_data, items in benchmark_config:
        for n_samples, dtype in items:
            time_results = []
            dist = {
                "float32":DistanceMetric32,
                "float64":DistanceMetric,
            }[dtype.__name__].get_metric("manhattan")

            for n_repeat in range(N_REPEATS):
                X, Y, _ = make_data(n_samples_X=n_samples, n_samples_Y=1, random_state=n_repeat)
                X = X.astype(dtype)
                start = perf_counter()
                dist.pairwise(X)
                duration = perf_counter() - start
                time_results.append(duration)
                writer.writerow(
                    {
                        "n_samples": n_samples,
                        "dtype": dtype.__name__,
                        "n_repeat": n_repeat,
                        "duration": duration,
                    }
                )
            results_mean, results_stdev = mean(time_results), stdev(time_results)
            print(
                f" {n_samples=} dtype={dtype.__name__} |"
                f" {results_mean:.3f} +/- {results_stdev:.3f}"
            )


import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
plt.rc('font', size=12)

_branches = ("main", "simd")
percentile_trim = .95
branches = {br:pd.read_csv(f'{results_path}{br}.csv') for br in _branches}

df = pd.concat([branches[br].assign(branch=br) for br in _branches])

group_by_attrs = ["dtype", "n_samples"]

grouped = list(df.groupby(group_by_attrs))
fig, axis = plt.subplots(2, 3, figsize=(14, 8), constrained_layout=True)
fig.patch.set_facecolor('white')
for (grouped_attrs, subset), ax in zip(grouped, axis.reshape(-1)):
    # Optionally trim outlier data
    if percentile_trim < 1:
        for branch in _branches:
            _subset = subset[subset["branch"]==branch]
            cut = _subset.duration < _subset.duration.quantile(percentile_trim)
            subset[subset["branch"]==branch] = _subset[cut]

    sns.violinplot(data=subset, y="duration", x="branch", ax=ax)
    ax.set_title("|".join( [f"{k}={v}" for k, v in zip(group_by_attrs, (grouped_attrs, ))] ))
    ax.set_xlabel("")
    
for ax in axis[:, 1:].ravel():
    ax.set_ylabel("")