From 95394d5bcbd7d73bae34c091a080c42bcfc9f07d Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Mon, 15 Nov 2021 21:40:37 +0100 Subject: [PATCH] Use set_options for asv bottleneck tests (#5986) * Use set_options for bottleneck tests * Use set_options in rolling * Update rolling.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update rolling.py * Update rolling.py * set_options not needed. Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- asv_bench/asv.conf.json | 2 +- asv_bench/benchmarks/dataarray_missing.py | 8 -- asv_bench/benchmarks/rolling.py | 92 ++++++++++++++--------- 3 files changed, 57 insertions(+), 45 deletions(-) diff --git a/asv_bench/asv.conf.json b/asv_bench/asv.conf.json index 83a2aa9f010..dafa0fc47e1 100644 --- a/asv_bench/asv.conf.json +++ b/asv_bench/asv.conf.json @@ -62,7 +62,7 @@ "pandas": [""], "netcdf4": [""], "scipy": [""], - "bottleneck": ["", null], + "bottleneck": [""], "dask": [""], "distributed": [""], }, diff --git a/asv_bench/benchmarks/dataarray_missing.py b/asv_bench/benchmarks/dataarray_missing.py index f89fe7f8eb9..d786c04e852 100644 --- a/asv_bench/benchmarks/dataarray_missing.py +++ b/asv_bench/benchmarks/dataarray_missing.py @@ -16,13 +16,6 @@ def make_bench_data(shape, frac_nan, chunks): return da -def requires_bottleneck(): - try: - import bottleneck # noqa: F401 - except ImportError: - raise NotImplementedError() - - class DataArrayMissingInterpolateNA: def setup(self, shape, chunks, limit): if chunks is not None: @@ -46,7 +39,6 @@ def time_interpolate_na(self, shape, chunks, limit): class DataArrayMissingBottleneck: def setup(self, shape, chunks, limit): - requires_bottleneck() if chunks is not None: requires_dask() self.da = make_bench_data(shape, 0.1, chunks) diff --git a/asv_bench/benchmarks/rolling.py b/asv_bench/benchmarks/rolling.py index f0e18bf2153..1d3713f19bf 100644 --- a/asv_bench/benchmarks/rolling.py +++ b/asv_bench/benchmarks/rolling.py @@ -36,29 +36,45 @@ def setup(self, *args, **kwargs): randn_long, dims="x", coords={"x": np.arange(long_nx) * 0.1} ) - @parameterized(["func", "center"], (["mean", "count"], [True, False])) - def time_rolling(self, func, center): - getattr(self.ds.rolling(x=window, center=center), func)().load() - - @parameterized(["func", "pandas"], (["mean", "count"], [True, False])) - def time_rolling_long(self, func, pandas): + @parameterized( + ["func", "center", "use_bottleneck"], + (["mean", "count"], [True, False], [True, False]), + ) + def time_rolling(self, func, center, use_bottleneck): + with xr.set_options(use_bottleneck=use_bottleneck): + getattr(self.ds.rolling(x=window, center=center), func)().load() + + @parameterized( + ["func", "pandas", "use_bottleneck"], + (["mean", "count"], [True, False], [True, False]), + ) + def time_rolling_long(self, func, pandas, use_bottleneck): if pandas: se = self.da_long.to_series() getattr(se.rolling(window=window, min_periods=window), func)() else: - getattr(self.da_long.rolling(x=window, min_periods=window), func)().load() - - @parameterized(["window_", "min_periods"], ([20, 40], [5, 5])) - def time_rolling_np(self, window_, min_periods): - self.ds.rolling(x=window_, center=False, min_periods=min_periods).reduce( - getattr(np, "nansum") - ).load() - - @parameterized(["center", "stride"], ([True, False], [1, 1])) - def time_rolling_construct(self, center, stride): - self.ds.rolling(x=window, center=center).construct( - "window_dim", stride=stride - ).sum(dim="window_dim").load() + with xr.set_options(use_bottleneck=use_bottleneck): + getattr( + self.da_long.rolling(x=window, min_periods=window), func + )().load() + + @parameterized( + ["window_", "min_periods", "use_bottleneck"], ([20, 40], [5, 5], [True, False]) + ) + def time_rolling_np(self, window_, min_periods, use_bottleneck): + with xr.set_options(use_bottleneck=use_bottleneck): + self.ds.rolling(x=window_, center=False, min_periods=min_periods).reduce( + getattr(np, "nansum") + ).load() + + @parameterized( + ["center", "stride", "use_bottleneck"], ([True, False], [1, 1], [True, False]) + ) + def time_rolling_construct(self, center, stride, use_bottleneck): + with xr.set_options(use_bottleneck=use_bottleneck): + self.ds.rolling(x=window, center=center).construct( + "window_dim", stride=stride + ).sum(dim="window_dim").load() class RollingDask(Rolling): @@ -87,24 +103,28 @@ def setup(self, *args, **kwargs): class DataArrayRollingMemory(RollingMemory): - @parameterized("func", ["sum", "max", "mean"]) - def peakmem_ndrolling_reduce(self, func): - roll = self.ds.var1.rolling(x=10, y=4) - getattr(roll, func)() + @parameterized(["func", "use_bottleneck"], (["sum", "max", "mean"], [True, False])) + def peakmem_ndrolling_reduce(self, func, use_bottleneck): + with xr.set_options(use_bottleneck=use_bottleneck): + roll = self.ds.var1.rolling(x=10, y=4) + getattr(roll, func)() - @parameterized("func", ["sum", "max", "mean"]) - def peakmem_1drolling_reduce(self, func): - roll = self.ds.var3.rolling(t=100) - getattr(roll, func)() + @parameterized(["func", "use_bottleneck"], (["sum", "max", "mean"], [True, False])) + def peakmem_1drolling_reduce(self, func, use_bottleneck): + with xr.set_options(use_bottleneck=use_bottleneck): + roll = self.ds.var3.rolling(t=100) + getattr(roll, func)() class DatasetRollingMemory(RollingMemory): - @parameterized("func", ["sum", "max", "mean"]) - def peakmem_ndrolling_reduce(self, func): - roll = self.ds.rolling(x=10, y=4) - getattr(roll, func)() - - @parameterized("func", ["sum", "max", "mean"]) - def peakmem_1drolling_reduce(self, func): - roll = self.ds.rolling(t=100) - getattr(roll, func)() + @parameterized(["func", "use_bottleneck"], (["sum", "max", "mean"], [True, False])) + def peakmem_ndrolling_reduce(self, func, use_bottleneck): + with xr.set_options(use_bottleneck=use_bottleneck): + roll = self.ds.rolling(x=10, y=4) + getattr(roll, func)() + + @parameterized(["func", "use_bottleneck"], (["sum", "max", "mean"], [True, False])) + def peakmem_1drolling_reduce(self, func, use_bottleneck): + with xr.set_options(use_bottleneck=use_bottleneck): + roll = self.ds.rolling(t=100) + getattr(roll, func)()