Skip to content

Commit

Permalink
add back the other asv tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Illviljan committed Nov 6, 2021
1 parent a89f62b commit 88a4a3d
Show file tree
Hide file tree
Showing 11 changed files with 1,062 additions and 0 deletions.
38 changes: 38 additions & 0 deletions asv_bench/benchmarks/combine.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
import numpy as np

import xarray as xr


class Combine:
"""Benchmark concatenating and merging large datasets"""

def setup(self):
"""Create 4 datasets with two different variables"""

t_size, x_size, y_size = 50, 450, 400
t = np.arange(t_size)
data = np.random.randn(t_size, x_size, y_size)

self.dsA0 = xr.Dataset(
{"A": xr.DataArray(data, coords={"T": t}, dims=("T", "X", "Y"))}
)
self.dsA1 = xr.Dataset(
{"A": xr.DataArray(data, coords={"T": t + t_size}, dims=("T", "X", "Y"))}
)
self.dsB0 = xr.Dataset(
{"B": xr.DataArray(data, coords={"T": t}, dims=("T", "X", "Y"))}
)
self.dsB1 = xr.Dataset(
{"B": xr.DataArray(data, coords={"T": t + t_size}, dims=("T", "X", "Y"))}
)

def time_combine_nested(self):
datasets = [[self.dsA0, self.dsA1], [self.dsB0, self.dsB1]]

xr.combine_nested(datasets, concat_dim=[None, "T"])

def time_combine_by_coords(self):
"""Also has to load and arrange t coordinate"""
datasets = [self.dsA0, self.dsA1, self.dsB0, self.dsB1]

xr.combine_by_coords(datasets)
80 changes: 80 additions & 0 deletions asv_bench/benchmarks/dataarray_missing.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
import pandas as pd

import xarray as xr

from . import parameterized, randn, requires_dask


def make_bench_data(shape, frac_nan, chunks):
vals = randn(shape, frac_nan)
coords = {"time": pd.date_range("2000-01-01", freq="D", periods=shape[0])}
da = xr.DataArray(vals, dims=("time", "x", "y"), coords=coords)

if chunks is not None:
da = da.chunk(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:
requires_dask()
self.da = make_bench_data(shape, 0.1, chunks)

@parameterized(
["shape", "chunks", "limit"],
(
[(365, 75, 75)],
[None, {"x": 25, "y": 25}],
[None, 3],
),
)
def time_interpolate_na(self, shape, chunks, limit):
actual = self.da.interpolate_na(dim="time", method="linear", limit=limit)

if chunks is not None:
actual = actual.compute()


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)

@parameterized(
["shape", "chunks", "limit"],
(
[(365, 75, 75)],
[None, {"x": 25, "y": 25}],
[None, 3],
),
)
def time_ffill(self, shape, chunks, limit):
actual = self.da.ffill(dim="time", limit=limit)

if chunks is not None:
actual = actual.compute()

@parameterized(
["shape", "chunks", "limit"],
(
[(365, 75, 75)],
[None, {"x": 25, "y": 25}],
[None, 3],
),
)
def time_bfill(self, shape, chunks, limit):
actual = self.da.ffill(dim="time", limit=limit)

if chunks is not None:
actual = actual.compute()
Loading

0 comments on commit 88a4a3d

Please sign in to comment.