Benchmarks for memview_to_ptr

bench.py

# %%
results_path = 'local_artifacts/benchmarks/dist'
results_path += '/' if results_path[-1] != '/' else ''

# %%
from sklearn.metrics._pairwise_distances_reduction import ArgKmin
from scipy.sparse import csr_matrix
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 = "PR"


METRIC = 'manhattan'
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

N_FEATURES = 10

benchmark_config = [
    (
        partial(_generate_PWD_data, n_features=N_FEATURES, n_classes=2),
        product(
            [2_000, 10_000],
            [2_000, 10_000],
            ["manhattan", "chebyshev", "minkowski"],
        ),
    ),
]


N_REPEATS = 10

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

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

    for make_data, items in benchmark_config:
        for n_samples, n_samples_test, metric 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_test, random_state=n_repeat)
                X = csr_matrix(X)
                Y = csr_matrix(Y)
                start = perf_counter()
                ArgKmin.compute(X, Y, 100, metric=metric)
                duration = perf_counter() - start
                time_results.append(duration)
                writer.writerow(
                    {
                        "n_samples": n_samples,
                        "n_samples_test": n_samples_test,
                        "metric": metric,
                        "n_repeat": n_repeat,
                        "duration": duration,
                    }
                )
            results_mean, results_stdev = mean(time_results), stdev(time_results)
            print(
                f" {n_samples=}, {n_samples_test=}, {metric=} |"
                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)

GRID_LAYOUT = (3, 4)
FIGURE_SIZE = (14, 12)

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"].mean()[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"].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, 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[:, 1:].ravel():
        ax.set_ylabel("")
    fig.suptitle(title, fontsize=18)
    plt.show()


# %%
_branches = ("main", "PR")
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 = ["metric", "n_samples_test", "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, default=_branches[0])

# generic_chart(_violen_perf, df.groupby(group_by_attrs), **default_args)
rel_title = f"ArgKmin.compute(X_csr, Y_csr, 100) relative performance (higher is better)\n"
abs_title = f"ArgKmin.compute(X_csr, Y_csr, 100) time spent (lower is better)\n"
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)