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)