|
| 1 | +from __future__ import absolute_import, division, print_function |
| 2 | + |
| 3 | +import numpy as np |
| 4 | +import pandas as pd |
| 5 | + |
| 6 | +import xarray as xr |
| 7 | + |
| 8 | +from . import parameterized, randn, requires_dask |
| 9 | + |
| 10 | +nx = 3000 |
| 11 | +long_nx = 30000000 |
| 12 | +ny = 2000 |
| 13 | +nt = 1000 |
| 14 | +window = 20 |
| 15 | + |
| 16 | +randn_xy = randn((nx, ny), frac_nan=0.1) |
| 17 | +randn_xt = randn((nx, nt)) |
| 18 | +randn_t = randn((nt, )) |
| 19 | +randn_long = randn((long_nx, ), frac_nan=0.1) |
| 20 | + |
| 21 | + |
| 22 | +new_x_short = np.linspace(0.3 * nx, 0.7 * nx, 100) |
| 23 | +new_x_long = np.linspace(0.3 * nx, 0.7 * nx, 1000) |
| 24 | +new_y_long = np.linspace(0.1, 0.9, 1000) |
| 25 | + |
| 26 | + |
| 27 | +class Interpolation(object): |
| 28 | + def setup(self, *args, **kwargs): |
| 29 | + self.ds = xr.Dataset( |
| 30 | + {'var1': (('x', 'y'), randn_xy), |
| 31 | + 'var2': (('x', 't'), randn_xt), |
| 32 | + 'var3': (('t', ), randn_t)}, |
| 33 | + coords={'x': np.arange(nx), |
| 34 | + 'y': np.linspace(0, 1, ny), |
| 35 | + 't': pd.date_range('1970-01-01', periods=nt, freq='D'), |
| 36 | + 'x_coords': ('x', np.linspace(1.1, 2.1, nx))}) |
| 37 | + |
| 38 | + @parameterized(['method', 'is_short'], |
| 39 | + (['linear', 'cubic'], [True, False])) |
| 40 | + def time_interpolation(self, method, is_short): |
| 41 | + new_x = new_x_short if is_short else new_x_long |
| 42 | + self.ds.interp(x=new_x, method=method).load() |
| 43 | + |
| 44 | + @parameterized(['method'], |
| 45 | + (['linear', 'nearest'])) |
| 46 | + def time_interpolation_2d(self, method): |
| 47 | + self.ds.interp(x=new_x_long, y=new_y_long, method=method).load() |
| 48 | + |
| 49 | + |
| 50 | +class InterpolationDask(Interpolation): |
| 51 | + def setup(self, *args, **kwargs): |
| 52 | + requires_dask() |
| 53 | + super(InterpolationDask, self).setup(**kwargs) |
| 54 | + self.ds = self.ds.chunk({'t': 50}) |
0 commit comments