Benchmarks for SIMD-accelerated distance metrics

bench.py

results_path = "local_artifacts/benchmarks/native/quick/manhattan"
results_path += "/" if results_path[-1] != "/" else ""

from distance_metrics import get_distance_metric

from sklearn.metrics._dist_metrics import ManhattanDistance, ManhattanDistance32
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

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 = 5

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": ManhattanDistance32(),
                "float64": ManhattanDistance(),
            }[dtype.__name__]
            # dist = get_distance_metric(np.array([0], dtype=dtype), '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
from textwrap import wrap

plt.rc("font", size=12)


def _violen_perf(subset, ax):
    sns.violinplot(data=subset, y="duration", x="branch", ax=ax)


def _rel_perf(subset, ax):
    base = subset.groupby("branch")["duration"].mean()["main"]
    subset["duration"] = base / subset["duration"]
    subset = subset.rename(columns={"duration": "relative performance"})
    graph = sns.barplot(
        subset, x="branch", y="relative performance", errorbar=None, ax=ax
    )
    graph.axhline(1, color="black")


def _abs_perf(subset, ax):
    base = subset.groupby("branch")["duration"].mean().min()
    subset = subset.rename(columns={"duration": "time (sec)"})
    graph = sns.barplot(subset, x="branch", y="time (sec)", errorbar=None, ax=ax)
    graph.axhline(base, color="black")


def generic_chart(func, grouped, percentile_trim, branches, group_by_attrs):
    grouped_list = list(grouped)
    fig, axis = plt.subplots(2, 3, figsize=(14, 9), constrained_layout=True)
    fig.patch.set_facecolor("white")
    for (grouped_attrs, subset), ax in zip(grouped_list, 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]

        func(subset, ax)

        ax.set_title(
            wrap(
                "|".join(
                    [f"{k}={v}" for k, v in zip(group_by_attrs, (grouped_attrs,))]
                ),
                40,
            )
        )
        ax.set_xlabel("")

    for ax in axis[:, 1:].ravel():
        ax.set_ylabel("")
    plt.show()
    
    
    
    
    

_branches = ("main", "xsimd")
percentile_trim = 0.9
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))
grouped_cp = list(df.groupby(group_by_attrs))

default_args = dict(
    percentile_trim=percentile_trim, branches=_branches, group_by_attrs=group_by_attrs
)

# generic_chart(_violen_perf, df.groupby(group_by_attrs), **default_args)
generic_chart(_rel_perf, df.groupby(group_by_attrs), **default_args)
generic_chart(_abs_perf, df.groupby(group_by_attrs), **default_args)