|
| 1 | +import numpy as np |
| 2 | + |
| 3 | +import xarray as xr |
| 4 | + |
| 5 | + |
| 6 | +class Combine: |
| 7 | + """Benchmark concatenating and merging large datasets""" |
| 8 | + |
| 9 | + def setup(self): |
| 10 | + """Create 4 datasets with two different variables""" |
| 11 | + |
| 12 | + t_size, x_size, y_size = 50, 450, 400 |
| 13 | + t = np.arange(t_size) |
| 14 | + data = np.random.randn(t_size, x_size, y_size) |
| 15 | + |
| 16 | + self.dsA0 = xr.Dataset( |
| 17 | + {"A": xr.DataArray(data, coords={"T": t}, dims=("T", "X", "Y"))} |
| 18 | + ) |
| 19 | + self.dsA1 = xr.Dataset( |
| 20 | + {"A": xr.DataArray(data, coords={"T": t + t_size}, dims=("T", "X", "Y"))} |
| 21 | + ) |
| 22 | + self.dsB0 = xr.Dataset( |
| 23 | + {"B": xr.DataArray(data, coords={"T": t}, dims=("T", "X", "Y"))} |
| 24 | + ) |
| 25 | + self.dsB1 = xr.Dataset( |
| 26 | + {"B": xr.DataArray(data, coords={"T": t + t_size}, dims=("T", "X", "Y"))} |
| 27 | + ) |
| 28 | + |
| 29 | + def time_combine_nested(self): |
| 30 | + datasets = [[self.dsA0, self.dsA1], [self.dsB0, self.dsB1]] |
| 31 | + |
| 32 | + xr.combine_nested(datasets, concat_dim=[None, "T"]) |
| 33 | + |
| 34 | + def time_combine_by_coords(self): |
| 35 | + """Also has to load and arrange t coordinate""" |
| 36 | + datasets = [self.dsA0, self.dsA1, self.dsB0, self.dsB1] |
| 37 | + |
| 38 | + xr.combine_by_coords(datasets) |
0 commit comments