def test_datetime_single_string():
da = xr.DataArray(np.arange(24), dims='time',
coords={'time': pd.date_range('2000-01-01', periods=24)})
actual = da.interp(time='2000-01-01T12:00')
expected = xr.DataArray(0.5)
assert_allclose(actual.drop('time'), expected)
def test_interpolate_1d(method, dim, case):
if not has_scipy:
pytest.skip('scipy is not installed.')
if not has_dask and case in [1]:
pytest.skip('dask is not installed in the environment.')
da = get_example_data(case)
xdest = np.linspace(0.0, 0.9, 80)
if dim == 'y' and case == 1:
with pytest.raises(NotImplementedError):
actual = da.interp(method=method, **{dim: xdest})
pytest.skip('interpolation along chunked dimension is '
'not yet supported')
actual = da.interp(method=method, **{dim: xdest})
# scipy interpolation for the reference
def func(obj, new_x):
return scipy.interpolate.interp1d(
da[dim], obj.data, axis=obj.get_axis_num(dim), bounds_error=False,
fill_value=np.nan, kind=method)(new_x)
if dim == 'x':
coords = {'x': xdest, 'y': da['y'], 'x2': ('x', func(da['x2'], xdest))}
else: # y
coords = {'x': da['x'], 'y': xdest, 'x2': da['x2']}
expected = xr.DataArray(func(da, xdest), dims=['x', 'y'], coords=coords)
assert_allclose(actual, expected)
def test_cftime():
times = xr.cftime_range('2000', periods=24, freq='D')
da = xr.DataArray(np.arange(24), coords=[times], dims='time')
times_new = xr.cftime_range('2000-01-01T12:00:00', periods=3, freq='D')
actual = da.interp(time=times_new)
expected = xr.DataArray([0.5, 1.5, 2.5], coords=[times_new], dims=['time'])
assert_allclose(actual, expected)
def test_datetime(x_new, expected):
da = xr.DataArray(np.arange(24), dims='time',
coords={'time': pd.date_range('2000-01-01', periods=24)})
actual = da.interp(time=x_new)
expected_da = xr.DataArray(np.atleast_1d(expected), dims=['time'],
coords={'time': (np.atleast_1d(x_new)
.astype('datetime64[ns]'))})
assert_allclose(actual, expected_da)
def test_interpolate_vectorize(use_dask):
if not has_scipy:
pytest.skip('scipy is not installed.')
if not has_dask and use_dask:
pytest.skip('dask is not installed in the environment.')
# scipy interpolation for the reference
def func(obj, dim, new_x):
shape = [s for i, s in enumerate(obj.shape)
if i != obj.get_axis_num(dim)]
for s in new_x.shape[::-1]:
shape.insert(obj.get_axis_num(dim), s)
return scipy.interpolate.interp1d(
da[dim], obj.data, axis=obj.get_axis_num(dim),
bounds_error=False, fill_value=np.nan)(new_x).reshape(shape)
da = get_example_data(0)
if use_dask:
da = da.chunk({'y': 5})
# xdest is 1d but has different dimension
xdest = xr.DataArray(np.linspace(0.1, 0.9, 30), dims='z',
coords={'z': np.random.randn(30),
'z2': ('z', np.random.randn(30))})
actual = da.interp(x=xdest, method='linear')
expected = xr.DataArray(func(da, 'x', xdest), dims=['z', 'y'],
coords={'z': xdest['z'], 'z2': xdest['z2'],
'y': da['y'],
'x': ('z', xdest.values),
'x2': ('z', func(da['x2'], 'x', xdest))})
assert_allclose(actual,
expected.transpose('z', 'y', transpose_coords=True))
# xdest is 2d
xdest = xr.DataArray(np.linspace(0.1, 0.9, 30).reshape(6, 5),
dims=['z', 'w'],
coords={'z': np.random.randn(6),
'w': np.random.randn(5),
'z2': ('z', np.random.randn(6))})
actual = da.interp(x=xdest, method='linear')
expected = xr.DataArray(
func(da, 'x', xdest),
dims=['z', 'w', 'y'],
coords={'z': xdest['z'], 'w': xdest['w'], 'z2': xdest['z2'],
'y': da['y'], 'x': (('z', 'w'), xdest),
'x2': (('z', 'w'), func(da['x2'], 'x', xdest))})
assert_allclose(actual,
expected.transpose('z', 'w', 'y', transpose_coords=True))
def test_cftime_single_string():
from cftime import DatetimeProlepticGregorian
times = xr.cftime_range('2000', periods=24, freq='D')
da = xr.DataArray(np.arange(24), coords=[times], dims='time')
times_new = '2000-01-01T12:00'
actual = da.interp(time=times_new)
times_new_array = _parse_array_of_cftime_strings(
np.array(times_new), DatetimeProlepticGregorian)
expected = xr.DataArray(0.5, coords={'time': times_new_array})
assert_allclose(actual, expected)
def test_rolling_exp_mean_pandas(self, da, dim, window_type, window) -> None:
da = da.isel(a=0).where(lambda x: x > 0.2)
result = da.rolling_exp(window_type=window_type, **{dim: window}).mean()
assert isinstance(result, DataArray)
pandas_array = da.to_pandas()
assert pandas_array.index.name == "time"
if dim == "x":
pandas_array = pandas_array.T
expected = xr.DataArray(
pandas_array.ewm(**{window_type: window}).mean()
).transpose(*da.dims)
assert_allclose(expected.variable, result.variable)
def test_crps_ensemble_accessor(o, f, dask_bool, outer_bool):
if dask_bool:
o = o.chunk()
f = f.chunk()
actual = crps_ensemble(o, f)
ds = xr.Dataset()
ds['o'] = o
ds['f'] = f
if outer_bool:
ds = ds.drop_vars('f')
expected = ds.xs.crps_ensemble('o', f)
else:
expected = ds.xs.crps_ensemble('o', 'f')
assert_allclose(actual, expected)
def test_cftime_list_of_strings():
from cftime import DatetimeProlepticGregorian
times = xr.cftime_range('2000', periods=24, freq='D')
da = xr.DataArray(np.arange(24), coords=[times], dims='time')
times_new = ['2000-01-01T12:00', '2000-01-02T12:00', '2000-01-03T12:00']
actual = da.interp(time=times_new)
times_new_array = _parse_array_of_cftime_strings(
np.array(times_new), DatetimeProlepticGregorian)
expected = xr.DataArray([0.5, 1.5, 2.5], coords=[times_new_array],
dims=['time'])
assert_allclose(actual, expected)
def test_datetime(x_new, expected):
da = xr.DataArray(
np.arange(24),
dims="time",
coords={"time": pd.date_range("2000-01-01", periods=24)},
)
actual = da.interp(time=x_new)
expected_da = xr.DataArray(
np.atleast_1d(expected),
dims=["time"],
coords={"time": (np.atleast_1d(x_new).astype("datetime64[ns]"))},
)
assert_allclose(actual, expected_da)
def test_rolling_reduce_nonnumeric(self, center, min_periods, window, name) -> None:
da = DataArray(
[0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time"
).isnull()
if min_periods is not None and window < min_periods:
min_periods = window
rolling_obj = da.rolling(time=window, center=center, min_periods=min_periods)
# add nan prefix to numpy methods to get similar behavior as bottleneck
actual = rolling_obj.reduce(getattr(np, "nan%s" % name))
expected = getattr(rolling_obj, name)()
assert_allclose(actual, expected)
assert actual.dims == expected.dims
def test_brier_score_accessor(o, f, threshold, dask_bool, outer_bool):
if dask_bool:
o = o.chunk()
f = f.chunk()
actual = brier_score(o > threshold, (f > threshold).mean('member'))
ds = xr.Dataset()
ds['o'] = o > threshold
ds['f'] = (f > threshold).mean('member')
if outer_bool:
ds = ds.drop_vars('f')
expected = ds.xs.brier_score('o', (f > threshold).mean('member'))
else:
expected = ds.xs.brier_score('o', 'f')
assert_allclose(actual, expected)
def test_pearson_r_p_value_xr(a, b, dim):
actual = xr.apply_ufunc(_pearson_r_p_value,
a,
b,
input_core_dims=[[dim], [dim]],
kwargs={'axis': -1},
dask='parallelized',
output_dtypes=[float])
_a = a.values
_b = b.values
axis = a.dims.index(dim)
res = _pearson_r_p_value(_a, _b, axis)
expected = actual.copy()
expected.values = res
assert_allclose(actual, expected)
def test_cftime_list_of_strings():
from cftime import DatetimeProlepticGregorian
times = xr.cftime_range('2000', periods=24, freq='D')
da = xr.DataArray(np.arange(24), coords=[times], dims='time')
times_new = ['2000-01-01T12:00', '2000-01-02T12:00', '2000-01-03T12:00']
actual = da.interp(time=times_new)
times_new_array = _parse_array_of_cftime_strings(
np.array(times_new), DatetimeProlepticGregorian)
expected = xr.DataArray([0.5, 1.5, 2.5], coords=[times_new_array],
dims=['time'])
assert_allclose(actual, expected)
def test_pearson_r_xr_dask(a_dask, b_dask, dim):
actual = xr.apply_ufunc(_pearson_r,
a_dask,
b_dask,
input_core_dims=[[dim], [dim]],
kwargs={'axis': -1},
dask='parallelized',
output_dtypes=[float])
_a_dask = a_dask.values
_b_dask = b_dask.values
axis = a_dask.dims.index(dim)
res = _pearson_r(_a_dask, _b_dask, axis)
expected = actual.copy()
expected.values = res
assert_allclose(actual, expected)
def test_crps_quadrature_accessor(o, dask_bool, outer_bool):
cdf_or_dist = norm
if dask_bool:
o = o.chunk()
actual = crps_quadrature(o, cdf_or_dist)
ds = xr.Dataset()
ds['o'] = o
ds['cdf_or_dist'] = cdf_or_dist
if outer_bool:
ds = ds.drop_vars('cdf_or_dist')
expected = ds.xs.crps_quadrature('o', cdf_or_dist)
else:
expected = ds.xs.crps_quadrature('o', 'cdf_or_dist')
assert_allclose(actual, expected)
def test_rolling_reduce(self, da, center, min_periods, window, name) -> None:
if min_periods is not None and window < min_periods:
min_periods = window
if da.isnull().sum() > 1 and window == 1:
# this causes all nan slices
window = 2
rolling_obj = da.rolling(time=window, center=center, min_periods=min_periods)
# add nan prefix to numpy methods to get similar # behavior as bottleneck
actual = rolling_obj.reduce(getattr(np, "nan%s" % name))
expected = getattr(rolling_obj, name)()
assert_allclose(actual, expected)
assert actual.dims == expected.dims
def test_threshold_brier_score_accessor(o, f, threshold, dask_bool,
outer_bool):
if dask_bool:
o = o.chunk()
f = f.chunk()
actual = threshold_brier_score(o, f, threshold)
ds = xr.Dataset()
ds['o'] = o
ds['f'] = f
if outer_bool:
ds = ds.drop_vars('f')
expected = ds.xs.threshold_brier_score('o', f, threshold)
else:
expected = ds.xs.threshold_brier_score('o', 'f', threshold)
assert_allclose(actual, expected)
def test_decompose(method: InterpOptions) -> None:
da = xr.DataArray(
np.arange(6).reshape(3, 2),
dims=["x", "y"],
coords={"x": [0, 1, 2], "y": [-0.1, -0.3]},
)
x_new = xr.DataArray([0.5, 1.5, 2.5], dims=["x1"])
y_new = xr.DataArray([-0.15, -0.25], dims=["y1"])
x_broadcast, y_broadcast = xr.broadcast(x_new, y_new)
assert x_broadcast.ndim == 2
actual = da.interp(x=x_new, y=y_new, method=method).drop_vars(("x", "y"))
expected = da.interp(x=x_broadcast, y=y_broadcast, method=method).drop_vars(
("x", "y")
)
assert_allclose(actual, expected)
def test_crps_quadrature_args(o_dask, f_prob_dask, keep_attrs):
"""Test crps_quadrature args."""
# to speed things up
o_dask = o_dask.isel(time=0, drop=True)
f_prob_dask = f_prob_dask.isel(time=0, drop=True)
xmin, xmax, tol = -10, 10, 1e-6
cdf_or_dist = norm
actual = crps_quadrature(
o_dask, cdf_or_dist, xmin, xmax, tol, keep_attrs=keep_attrs
)
expected = properscoring.crps_quadrature(o_dask, cdf_or_dist, xmin, xmax, tol)
expected = xr.DataArray(expected, coords=o_dask.coords).mean()
# test for numerical identity of xskillscore crps and crps
assert_allclose(actual, expected)
# test that xskillscore crps_ensemble returns chunks
assert_chunk(actual, True)
def test_skipna_returns_same_value_as_dropped_pairwise_nans_with_weights(
a_1d_nan, b_1d_nan, weights_time, metric):
"""Tests that DataArrays with pairwise nans return the same result
as the same two with those nans dropped."""
a_dropped, b_dropped, weights_time_dropped = drop_nans(
a_1d_nan, b_1d_nan, weights_time)
res_with_nans = metric(a_1d_nan,
b_1d_nan,
"time",
skipna=True,
weights=weights_time)
res_dropped_nans = metric(a_dropped,
b_dropped,
"time",
weights=weights_time_dropped)
assert_allclose(res_with_nans, res_dropped_nans)
def test_crps_ensemble_api_and_inputs(o, f_prob, keep_attrs, input_type, chunk_bool):
"""Test that crps_ensemble keeps attributes, chunking, input types and equals
properscoring.crps_ensemble."""
o, f_prob = modify_inputs(o, f_prob, input_type, chunk_bool)
actual = crps_ensemble(o, f_prob, keep_attrs=keep_attrs)
if input_type == "DataArray": # properscoring allows only DataArrays
expected = properscoring.crps_ensemble(o, f_prob, axis=0)
expected = xr.DataArray(expected, coords=o.coords).mean()
# test for numerical identity of xskillscore crps and properscoring crps
assert_allclose(actual, expected)
# test that returns chunks
assert_chunk(actual, chunk_bool)
# test that attributes are kept
assert_keep_attrs(actual, o, keep_attrs)
# test that input types equal output types
assign_type_input_output(actual, o)
def test_rps_new_identical_old_xhistogram(o, f_prob, fair_bool):
"""Test that new rps algorithm is identical to old algorithm with xhistogram.
Makes a difference whether full range of f_prob is covered or not."""
category_edges_np = np.array([0, 0.2, 0.4, 0.6, 0.8, 1.0])
category_edges_xr = xr.DataArray(
category_edges_np,
dims="category_edge",
coords={"category_edge": category_edges_np},
)
dim = "time"
actual = rps(o, f_prob, dim=dim, category_edges=category_edges_xr)
expected = rps_xhist(o, f_prob, dim=dim, category_edges=category_edges_np)
drop = ["observations_category_edge", "forecasts_category_edge"]
assert_allclose(
actual.rename("histogram_category_edge").drop(drop),
expected.drop(drop))
def test_correlation_metrics_ufunc_same_np(a, b, dim, weight_bool,
weights_cos_lat, metrics, skipna,
has_nan):
"""Test whether correlation metric for xarray functions (from
deterministic.py) give save numerical results as for numpy functions (from
np_deterministic.py)."""
if has_nan:
a[0] = np.nan
# unpack metrics
metric, _metric = metrics
# Only apply over time dimension for effective p value.
if (dim != "time") and (metric in temporal_only_metrics):
dim = "time"
# Generates subsetted weights to pass in as arg to main function and for
# the numpy testing.
_weights = adjust_weights(dim, weight_bool, weights_cos_lat)
if metric in temporal_only_metrics:
actual = metric(a, b, dim, skipna=skipna)
else:
actual = metric(a, b, dim, weights=_weights, skipna=skipna)
# check that no chunks for no chunk inputs
assert actual.chunks is None
dim, _ = _preprocess_dims(dim, a)
if len(dim) > 1:
new_dim = "_".join(dim)
_a = a.stack(**{new_dim: dim})
_b = b.stack(**{new_dim: dim})
if weight_bool:
_weights = _weights.stack(**{new_dim: dim})
else:
new_dim = dim[0]
_a = a
_b = b
_weights = _preprocess_weights(_a, dim, new_dim, _weights)
# ensure _weights.values or None
_weights = None if _weights is None else _weights.values
axis = _a.dims.index(new_dim)
if metric in temporal_only_metrics:
res = _metric(_a.values, _b.values, axis, skipna=skipna)
else:
res = _metric(_a.values, _b.values, _weights, axis, skipna=skipna)
expected = actual.copy()
expected.values = res
assert_allclose(actual, expected)
def test_cftime_list_of_strings() -> None:
from cftime import DatetimeProlepticGregorian
times = xr.cftime_range(
"2000", periods=24, freq="D", calendar="proleptic_gregorian"
)
da = xr.DataArray(np.arange(24), coords=[times], dims="time")
times_new = ["2000-01-01T12:00", "2000-01-02T12:00", "2000-01-03T12:00"]
actual = da.interp(time=times_new)
times_new_array = _parse_array_of_cftime_strings(
np.array(times_new), DatetimeProlepticGregorian
)
expected = xr.DataArray([0.5, 1.5, 2.5], coords=[times_new_array], dims=["time"])
assert_allclose(actual, expected)
def test_weighted_operations_nonequal_coords():
# There are no weights for a == 4, so that data point is ignored.
weights = DataArray(np.random.randn(4), dims=("a",), coords=dict(a=[0, 1, 2, 3]))
data = DataArray(np.random.randn(4), dims=("a",), coords=dict(a=[1, 2, 3, 4]))
check_weighted_operations(data, weights, dim="a", skipna=None)
q = 0.5
result = data.weighted(weights).quantile(q, dim="a")
# Expected value computed using code from https://aakinshin.net/posts/weighted-quantiles/ with values at a=1,2,3
expected = DataArray([0.9308707], coords={"quantile": [q]}).squeeze()
assert_allclose(result, expected)
data = data.to_dataset(name="data")
check_weighted_operations(data, weights, dim="a", skipna=None)
result = data.weighted(weights).quantile(q, dim="a")
assert_allclose(result, expected.to_dataset(name="data"))
def test_cftime_single_string():
from cftime import DatetimeProlepticGregorian
times = xr.cftime_range("2000",
periods=24,
freq="D",
calendar="proleptic_gregorian")
da = xr.DataArray(np.arange(24), coords=[times], dims="time")
times_new = "2000-01-01T12:00"
actual = da.interp(time=times_new)
times_new_array = _parse_array_of_cftime_strings(
np.array(times_new), DatetimeProlepticGregorian)
expected = xr.DataArray(0.5, coords={"time": times_new_array})
assert_allclose(actual, expected)
def test_mae_r_xr(a, b, dim, weight, weights_ones, weights_latitude):
# Generates subsetted weights to pass in as arg to main function and for the numpy testing.
weights_arg, weights_np = adjust_weights(weight, dim, weights_ones,
weights_latitude)
actual = mae(a, b, dim, weights=weights_arg)
assert actual.chunks is None
dim, axis = _preprocess_dims(dim)
_a = a
_b = b
_weights_np = _preprocess_weights(_a, dim, dim, weights_np)
axis = tuple(a.dims.index(d) for d in dim)
res = _mae(_a.values, _b.values, _weights_np.values, axis)
expected = actual.copy()
expected.values = res
assert_allclose(actual, expected)
def test_crps_quadrature_dask(o_dask, keep_attrs):
# to speed things up
o_dask = o_dask.isel(time=0, drop=True)
cdf_or_dist = norm
actual = crps_quadrature(o_dask, cdf_or_dist, keep_attrs=keep_attrs)
expected = properscoring.crps_quadrature(o_dask, cdf_or_dist)
expected = xr.DataArray(expected, coords=o_dask.coords).mean()
# test for numerical identity of xskillscore crps and properscoring crps
assert_allclose(actual, expected)
# test that xskillscore crps_ensemble returns chunks
assert actual.chunks is not None
# show that properscoring crps_ensemble returns no chunks
assert expected.chunks is None
if keep_attrs:
assert actual.attrs == o_dask.attrs
else:
assert actual.attrs == {}
def test_pearson_r_xr(a, b, dim):
actual = pearson_r(a, b, dim)
dim, _ = _preprocess(dim)
if len(dim) > 1:
new_dim = '_'.join(dim)
_a = a.stack(**{new_dim: dim})
_b = b.stack(**{new_dim: dim})
else:
new_dim = dim[0]
_a = a
_b = b
axis = _a.dims.index(new_dim)
res = _pearson_r(_a.values, _b.values, axis)
expected = actual.copy()
expected.values = res
assert_allclose(actual, expected)
def test_pearson_r_xr_dask(a_dask, b_dask, dim):
actual = pearson_r(a_dask, b_dask, dim)
dim, _ = _preprocess(dim)
if len(dim) > 1:
new_dim = '_'.join(dim)
_a_dask = a_dask.stack(**{new_dim: dim})
_b_dask = b_dask.stack(**{new_dim: dim})
else:
new_dim = dim[0]
_a_dask = a_dask
_b_dask = b_dask
axis = _a_dask.dims.index(new_dim)
res = _pearson_r(_a_dask.values, _b_dask.values, axis)
expected = actual.copy()
expected.values = res
assert_allclose(actual, expected)
def test_crps_gaussian_accessor(o, f, dask_bool, outer_bool):
if dask_bool:
o = o.chunk()
f = f.chunk()
mu = f.mean('member')
sig = f.std('member')
actual = crps_gaussian(o, mu, sig)
ds = xr.Dataset()
ds['o'] = o
ds['mu'] = mu
ds['sig'] = sig
if outer_bool:
ds = ds.drop_vars('mu')
expected = ds.xs.crps_gaussian('o', mu, sig)
else:
expected = ds.xs.crps_gaussian('o', 'mu', 'sig')
assert_allclose(actual, expected)
def test_ndrolling_construct(self, center, fill_value) -> None:
da = DataArray(
np.arange(5 * 6 * 7).reshape(5, 6, 7).astype(float),
dims=["x", "y", "z"],
coords={"x": ["a", "b", "c", "d", "e"], "y": np.arange(6)},
)
actual = da.rolling(x=3, z=2, center=center).construct(
x="x1", z="z1", fill_value=fill_value
)
if not isinstance(center, tuple):
center = (center, center)
expected = (
da.rolling(x=3, center=center[0])
.construct(x="x1", fill_value=fill_value)
.rolling(z=2, center=center[1])
.construct(z="z1", fill_value=fill_value)
)
assert_allclose(actual, expected)
def test_mse_r_xr_dask(a_dask, b_dask, dim, weight, weights_ones_dask,
weights_latitude_dask):
# Generates subsetted weights to pass in as arg to main function and for the numpy testing.
weights_arg, weights_np = adjust_weights(weight, dim, weights_ones_dask,
weights_latitude_dask)
actual = mse(a_dask, b_dask, dim, weights=weights_arg)
assert actual.chunks is not None
dim, axis = _preprocess_dims(dim)
_a_dask = a_dask
_b_dask = b_dask
_weights_np = _preprocess_weights(_a_dask, dim, dim, weights_np)
axis = tuple(a_dask.dims.index(d) for d in dim)
res = _mse(_a_dask.values, _b_dask.values, _weights_np.values, axis)
expected = actual.copy()
expected.values = res
assert_allclose(actual, expected)
def test_rank_histogram_sum(o, f_prob, dim, obj):
"""Test that the number of samples in the rank histogram is correct"""
if "ds" in obj:
name = "var"
o = o.to_dataset(name=name)
f_prob = f_prob.to_dataset(name=name)
if "chunked" in obj:
o = o.chunk()
f_prob = f_prob.chunk()
if dim == []:
with pytest.raises(ValueError):
rank_histogram(o, f_prob, dim=dim)
else:
rank_hist = rank_histogram(o, f_prob, dim=dim)
if "ds" in obj:
rank_hist = rank_hist[name]
o = o[name]
assert_allclose(rank_hist.sum(), o.count())
def test_interpolate_1d_methods(method):
if not has_scipy:
pytest.skip('scipy is not installed.')
da = get_example_data(0)
dim = 'x'
xdest = np.linspace(0.0, 0.9, 80)
actual = da.interp(method=method, **{dim: xdest})
# scipy interpolation for the reference
def func(obj, new_x):
return scipy.interpolate.interp1d(
da[dim], obj.data, axis=obj.get_axis_num(dim), bounds_error=False,
fill_value=np.nan, kind=method)(new_x)
coords = {'x': xdest, 'y': da['y'], 'x2': ('x', func(da['x2'], xdest))}
expected = xr.DataArray(func(da, xdest), dims=['x', 'y'], coords=coords)
assert_allclose(actual, expected)
def test_rank_histogram_sum(o, f_prob, dim, chunk_bool, input_type):
"""Test that the number of samples in the rank histogram is correct"""
o, f_prob = modify_inputs(o, f_prob, input_type, chunk_bool)
if dim == []:
with pytest.raises(ValueError):
rank_histogram(o, f_prob, dim=dim)
else:
rank_hist = rank_histogram(o, f_prob, dim=dim)
if "Dataset" in input_type:
rank_hist = rank_hist[list(o.data_vars)[0]]
o = o[list(o.data_vars)[0]]
assert_allclose(rank_hist.sum(), o.count())
assert_allclose(rank_hist.sum(), o.count())
# test that returns chunks
assert_chunk(rank_hist, chunk_bool)
# test that attributes are kept # TODO: add
# assert_keep_attrs(rank_hist, o, keep_attrs)
# test that input types equal output types
assign_type_input_output(rank_hist, o)
def test_dataset():
ds = create_test_data()
ds.attrs['foo'] = 'var'
ds['var1'].attrs['buz'] = 'var2'
new_dim2 = xr.DataArray([0.11, 0.21, 0.31], dims='z')
interpolated = ds.interp(dim2=new_dim2)
assert_allclose(interpolated['var1'], ds['var1'].interp(dim2=new_dim2))
assert interpolated['var3'].equals(ds['var3'])
# make sure modifying interpolated does not affect the original dataset
interpolated['var1'][:, 1] = 1.0
interpolated['var2'][:, 1] = 1.0
interpolated['var3'][:, 1] = 1.0
assert not interpolated['var1'].equals(ds['var1'])
assert not interpolated['var2'].equals(ds['var2'])
assert not interpolated['var3'].equals(ds['var3'])
# attrs should be kept
assert interpolated.attrs['foo'] == 'var'
assert interpolated['var1'].attrs['buz'] == 'var2'
def test_interpolate_scalar(method, case):
if not has_scipy:
pytest.skip('scipy is not installed.')
if not has_dask and case in [1]:
pytest.skip('dask is not installed in the environment.')
da = get_example_data(case)
xdest = 0.4
actual = da.interp(x=xdest, method=method)
# scipy interpolation for the reference
def func(obj, new_x):
return scipy.interpolate.interp1d(
da['x'], obj.data, axis=obj.get_axis_num('x'), bounds_error=False,
fill_value=np.nan)(new_x)
coords = {'x': xdest, 'y': da['y'], 'x2': func(da['x2'], xdest)}
expected = xr.DataArray(func(da, xdest), dims=['y'], coords=coords)
assert_allclose(actual, expected)
def test_interp_like():
ds = create_test_data()
ds.attrs['foo'] = 'var'
ds['var1'].attrs['buz'] = 'var2'
other = xr.DataArray(np.random.randn(3), dims=['dim2'],
coords={'dim2': [0, 1, 2]})
interpolated = ds.interp_like(other)
assert_allclose(interpolated['var1'],
ds['var1'].interp(dim2=other['dim2']))
assert_allclose(interpolated['var1'],
ds['var1'].interp_like(other))
assert interpolated['var3'].equals(ds['var3'])
# attrs should be kept
assert interpolated.attrs['foo'] == 'var'
assert interpolated['var1'].attrs['buz'] == 'var2'
other = xr.DataArray(np.random.randn(3), dims=['dim3'],
coords={'dim3': ['a', 'b', 'c']})
actual = ds.interp_like(other)
expected = ds.reindex_like(other)
assert_allclose(actual, expected)
def test_interpolate_nd_scalar(method, case):
if not has_scipy:
pytest.skip('scipy is not installed.')
if not has_dask and case in [4]:
pytest.skip('dask is not installed in the environment.')
da = get_example_data(case)
xdest = 0.4
ydest = 0.05
actual = da.interp(x=xdest, y=ydest, method=method)
# scipy interpolation for the reference
expected_data = scipy.interpolate.RegularGridInterpolator(
(da['x'], da['y']), da.transpose('x', 'y', 'z').values,
method='linear', bounds_error=False,
fill_value=np.nan)(np.stack([xdest, ydest], axis=-1))
coords = {'x': xdest, 'y': ydest, 'x2': da['x2'].interp(x=xdest),
'z': da['z']}
expected = xr.DataArray(expected_data[0], dims=['z'], coords=coords)
assert_allclose(actual, expected)
def test_sorted():
# unsorted non-uniform gridded data
x = np.random.randn(100)
y = np.random.randn(30)
z = np.linspace(0.1, 0.2, 10) * 3.0
da = xr.DataArray(
np.cos(x[:, np.newaxis, np.newaxis]) * np.cos(
y[:, np.newaxis]) * z,
dims=['x', 'y', 'z'],
coords={'x': x, 'y': y, 'x2': ('x', x**2), 'z': z})
x_new = np.linspace(0, 1, 30)
y_new = np.linspace(0, 1, 20)
da_sorted = da.sortby('x')
assert_allclose(da.interp(x=x_new),
da_sorted.interp(x=x_new, assume_sorted=True))
da_sorted = da.sortby(['x', 'y'])
assert_allclose(da.interp(x=x_new, y=y_new),
da_sorted.interp(x=x_new, y=y_new, assume_sorted=True))
with pytest.raises(ValueError):
da.interp(x=[0, 1, 2], assume_sorted=True)
def test_datetime():
da = xr.DataArray(np.random.randn(24), dims='time',
coords={'time': pd.date_range('2000-01-01', periods=24)})
x_new = pd.date_range('2000-01-02', periods=3)
actual = da.interp(time=x_new)
expected = da.isel(time=[1, 2, 3])
assert_allclose(actual, expected)
x_new = np.array([np.datetime64('2000-01-01T12:00'),
np.datetime64('2000-01-02T12:00')])
actual = da.interp(time=x_new)
assert_allclose(actual.isel(time=0).drop('time'),
0.5 * (da.isel(time=0) + da.isel(time=1)))
assert_allclose(actual.isel(time=1).drop('time'),
0.5 * (da.isel(time=1) + da.isel(time=2)))
def test_interpolate_nd(case):
if not has_scipy:
pytest.skip('scipy is not installed.')
if not has_dask and case == 4:
pytest.skip('dask is not installed in the environment.')
da = get_example_data(case)
# grid -> grid
xdest = np.linspace(0.1, 1.0, 11)
ydest = np.linspace(0.0, 0.2, 10)
actual = da.interp(x=xdest, y=ydest, method='linear')
# linear interpolation is separateable
expected = da.interp(x=xdest, method='linear')
expected = expected.interp(y=ydest, method='linear')
assert_allclose(actual.transpose('x', 'y', 'z'),
expected.transpose('x', 'y', 'z'))
# grid -> 1d-sample
xdest = xr.DataArray(np.linspace(0.1, 1.0, 11), dims='y')
ydest = xr.DataArray(np.linspace(0.0, 0.2, 11), dims='y')
actual = da.interp(x=xdest, y=ydest, method='linear')
# linear interpolation is separateable
expected_data = scipy.interpolate.RegularGridInterpolator(
(da['x'], da['y']), da.transpose('x', 'y', 'z').values,
method='linear', bounds_error=False,
fill_value=np.nan)(np.stack([xdest, ydest], axis=-1))
expected = xr.DataArray(
expected_data, dims=['y', 'z'],
coords={'z': da['z'], 'y': ydest, 'x': ('y', xdest.values),
'x2': da['x2'].interp(x=xdest)})
assert_allclose(actual.transpose('y', 'z'), expected)
# reversed order
actual = da.interp(y=ydest, x=xdest, method='linear')
assert_allclose(actual.transpose('y', 'z'), expected)
