SIMD Benchmark

bench.py

# %%
results_path = 'local_artifacts/benchmarks/simd'
results_path += '/' if results_path[-1] != '/' else ''
branch = "PR"

# %%
from sklearn.metrics import DistanceMetric
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)

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_X,) if n_outs == 1 else (n_samples_X, n_outs)
    y = rng.randint(n_classes, size=y_shape)

    return X, Y, y

benchmark_config = [
    (
        partial(_generate_PWD_data, n_classes=2),
        product(
            [2_000, 5_000],
            [10, 50, 150],
            [np.float32, np.float64],
        ),
    ),
]


N_REPEATS = 50
METRIC="manhattan"
with open(f'{results_path}{branch}.csv', 'w', newline='') as csvfile:

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

    for make_data, items in benchmark_config:
        for n_samples, n_features, dtype in items:
            time_results = []

            for n_repeat in range(N_REPEATS):
                X, Y, y = make_data(n_samples_X=n_samples, n_samples_Y=n_samples, n_features=n_features, random_state=n_repeat)
                X = X.astype(dtype)
                dst = DistanceMetric.get_metric(METRIC, dtype=dtype)
                start = perf_counter()
                dst.pairwise(X)
                duration = perf_counter() - start
                time_results.append(duration)
                writer.writerow(
                    {
                        "n_samples": n_samples,
                        "n_features": n_features,
                        "dtype": dtype.__name__,
                        "n_repeat": n_repeat,
                        "duration": duration,
                    }
                )
            results_min = min(time_results)
            print(
                f" {n_samples=}, {n_features=}, dtype={dtype.__name__} |"
                f" {results_min:.3f}"
            )

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

GRID_LAYOUT = (2, 6)
FIGURE_SIZE = (12, 8)

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

def _rel_perf(subset, ax, default, **kwargs):
    base = subset.groupby("branch")["duration"].min()[default]
    subset["duration"] = base / subset["duration"]
    y_title = "speedup vs main"
    subset = subset.rename(columns={"duration":y_title})
    graph = sns.barplot(subset, x="branch", y=y_title, errorbar=None, ax=ax)
    graph.axhline(1, color="black")

def _abs_perf(subset, ax, **kwargs):
    base = subset.groupby("branch")["duration"].min().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, title, **kwargs):
    grouped_list = list(grouped)
    fig, axis = plt.subplots(*GRID_LAYOUT, figsize=FIGURE_SIZE, 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, **kwargs)

        ax.set_title("\n".join( [f"{k}={v}" for k, v in zip(group_by_attrs, grouped_attrs)]))
        ax.set_xlabel("")
        
    for ax in axis.ravel():
        ax.set_ylabel("")
        ax.tick_params(axis='x', rotation=80)

    fig.suptitle(title, fontsize=18)
    plt.show()


# %%
_branches = ("main", "PR", "PR_noruntime", "PR_nobuild")
percentile_trim = .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 = ["n_samples", "n_features", "dtype"]
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, default=_branches[0])

# generic_chart(_violen_perf, df.groupby(group_by_attrs), **default_args)
rel_title = f"ManhattanDistance.pairwise() relative performance (higher is better)\n"
abs_title = f"ManhattanDistance.pairwise() time spent (lower is better)\n"
# generic_chart(_violen_perf, df.groupby(group_by_attrs), title=rel_title, **default_args)
generic_chart(_rel_perf, df.groupby(group_by_attrs), title=rel_title, **default_args)
generic_chart(_abs_perf, df.groupby(group_by_attrs), title=abs_title, **default_args)