Small ASV benchmark for scikit-learn/scikit-learn#19418

sklearn_19418_bench.py

from sklearn.feature_selection import mutual_info_regression, mutual_info_classif
from sklearn.neighbors import KernelDensity, NearestNeighbors
from .common import Benchmark
from sklearn.datasets import make_classification, make_regression


class RemovedCheckBenchmarks(Benchmark):

    param_names = ['n', 'd']
    params = (
        [100, 1000, 10_000],
        [10, 100],
    )


class MutualInfoCCRemovedCheck(RemovedCheckBenchmarks):

    def setup(self, n, d):
        self.X, self.y = make_regression(n_samples=n, n_features=d)

    def time_mi_cc(self, n, d):
        mutual_info_regression(self.X, self.y, discrete_features=False)


class MutualInfoCDRemovedCheck(RemovedCheckBenchmarks):

    def setup(self, n, d):
        self.X, self.y = make_classification(n_samples=n, n_features=d)

    def time_mi_cd(self, n, d):
        mutual_info_classif(self.X, self.y, discrete_features=False)


class KernelDensityRemovedCheck(RemovedCheckBenchmarks):

    def setup(self, n, d):
        self.X, y = make_regression(n_samples=n, n_features=d)
        self.kde = KernelDensity().fit(self.X, y)

    def time_query(self, n, d):
        self.kde.score_samples(self.X)


class NeighborsBaseCreationRemovedCheck(RemovedCheckBenchmarks):

    def setup(self, n, d):
        self.nn = NearestNeighbors(algorithm='kd_tree')
        self.X, _ = make_regression(n_samples=n, n_features=d)

    def time_creation(self, n, d):
        self.nn.fit(self.X)


class NeighborsBaseQueryRemovedCheck(RemovedCheckBenchmarks):

    def setup(self, n, d):
        self.X, _ = make_regression(n_samples=n, n_features=d)
        self.nn = NearestNeighbors(algorithm='kd_tree').fit(self.X)

    def time_radius_neighbors(self, n, d):
        self.nn.radius_neighbors(self.X)

    def time_kneighbors(self, n, d):
        self.nn.kneighbors(self.X)