Skip to content

Commit

Permalink
Use set_options for asv bottleneck tests (pydata#5986)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
Illviljan and pre-commit-ci[bot] authored Nov 15, 2021
1 parent b2d7cd8 commit 95394d5
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 45 deletions.
2 changes: 1 addition & 1 deletion asv_bench/asv.conf.json
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@
"pandas": [""],
"netcdf4": [""],
"scipy": [""],
"bottleneck": ["", null],
"bottleneck": [""],
"dask": [""],
"distributed": [""],
},
Expand Down
8 changes: 0 additions & 8 deletions asv_bench/benchmarks/dataarray_missing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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)
Expand Down
92 changes: 56 additions & 36 deletions asv_bench/benchmarks/rolling.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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)()

0 comments on commit 95394d5

Please sign in to comment.