def test_pearson_r_nd(a, b):
axis = 0
expected = np.squeeze(a[0, :, :]).copy()
for i in range(np.shape(a)[1]):
for j in range(np.shape(a)[2]):
_a = a[:, i, j]
_b = b[:, i, j]
expected[i, j], p = stats.pearsonr(_a, _b)
actual = _pearson_r(a, b, weights, axis, skipna)
assert_allclose(actual, expected)
axis = 1
expected = np.squeeze(a[:, 0, :]).copy()
for i in range(np.shape(a)[0]):
for j in range(np.shape(a)[2]):
_a = a[i, :, j]
_b = b[i, :, j]
expected[i, j], p = stats.pearsonr(_a, _b)
actual = _pearson_r(a, b, weights, axis, skipna)
assert_allclose(actual, expected)
axis = 2
expected = np.squeeze(a[:, :, 0]).copy()
for i in range(np.shape(a)[0]):
for j in range(np.shape(a)[1]):
_a = a[i, j, :]
_b = b[i, j, :]
expected[i, j], p = stats.pearsonr(_a, _b)
actual = _pearson_r(a, b, weights, axis, skipna)
assert_allclose(actual, expected)
def test_pearson_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 = pearson_r(a_dask, b_dask, dim, weights=weights_arg)
assert actual.chunks is not None
dim, _ = _preprocess_dims(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})
_weights_np = weights_np.stack(**{new_dim: dim})
else:
new_dim = dim[0]
_a_dask = a_dask
_b_dask = b_dask
_weights_np = weights_np
_weights_np = _preprocess_weights(_a_dask, dim, new_dim, _weights_np)
axis = _a_dask.dims.index(new_dim)
res = _pearson_r(_a_dask.values, _b_dask.values, _weights_np.values, axis)
expected = actual.copy()
expected.values = res
assert_allclose(actual, expected)
def test_pearson_r_xr(a, b, dim):
actual = xr.apply_ufunc(_pearson_r,
a,
b,
input_core_dims=[[dim], [dim]],
kwargs={'axis': -1})
_a = a.values
_b = b.values
axis = a.dims.index(dim)
res = _pearson_r(_a, _b, axis)
expected = actual.copy()
expected.values = res
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_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_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)